Merge K Sorted Arrays

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

public class Solution {
    public int[] mergeKSortedArrays(int[][] arrays) {
        int k = arrays.length;
        if (k == 0) return new int[0];
        return mergeKArrays(arrays, 0, k - 1);
    }

    private int[] mergeKArrays(int[][] arrays, int start, int end) {
        if (start == end) {
            return arrays[start];
        }

        int mid = start + (end - start) / 2;
        int[] left = mergeKArrays(arrays, start, mid);
        int[] right = mergeKArrays(arrays, mid + 1, end);

        return mergeTwoArrays(left, right);
    }

    private int[] mergeTwoArrays(int[] arr1, int[] arr2) {
        int len1 = arr1.length, len2 = arr2.length;
        int[] result = new int[len1 + len2];
        int i = 0, j = 0, k = 0;

        while (i < len1 && j < len2) {
            if (arr1[i] <= arr2[j]) {
                result[k++] = arr1[i++];
            } else {
                result[k++] = arr2[j++];
            }
        }

        while (i < len1) {
            result[k++] = arr1[i++];
        }

        while (j < len2) {
            result[k++] = arr2[j++];
        }

        return result;
    }

    public static void main(String[] args) {
        Solution sol = new Solution();
        int[][] arrays = {
            {1, 4, 7},
            {2, 5, 8},
            {3, 6, 9}
        };
        
        int[] result = sol.mergeKSortedArrays(arrays);
        System.out.println("Merged array: " + Arrays.toString(result));  // Expected output: [1, 2, 3, 4, 5, 6, 7, 8, 9]
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

class Solution:
    def mergeKSortedArrays(self, arrays: List[List[int]]) -> List[int]:
        if not arrays:
            return []
        return self.mergeKArrays(arrays, 0, len(arrays) - 1)

    def mergeKArrays(self, arrays: List[List[int]], start: int, end: int) -> List[int]:
        if start == end:
            return arrays[start]

        mid = start + (end - start) // 2
        left = self.mergeKArrays(arrays, start, mid)
        right = self.mergeKArrays(arrays, mid + 1, end)

        return self.mergeTwoArrays(left, right)

    def mergeTwoArrays(self, arr1: List[int], arr2: List[int]) -> List[int]:
        len1, len2 = len(arr1), len(arr2)
        result = []
        i = j = 0

        while i < len1 and j < len2:
            if arr1[i] <= arr2[j]:
                result.append(arr1[i])
                i += 1
            else:
                result.append(arr2[j])
                j += 1

        while i < len1:
            result.append(arr1[i])
            i += 1

        while j < len2:
        result.append(arr2[j])
        j += 1

        return result

# Example usage
sol = Solution()
arrays = [
    [1, 4, 7],
    [2, 5, 8],
    [3, 6, 9]
]
result = sol.mergeKSortedArrays(arrays)
print("Merged array:", result)  # Expected output: [1, 2, 3, 4, 5, 6, 7, 8, 9]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

public class Solution {
    public int[] mergeKSortedArrays(int[][] arrays) {
        PriorityQueue<Element> minHeap = new PriorityQueue<>((a, b) -> a.value - b.value);
        int totalSize = 0;
        
        // Initialize the heap with the first element of each array
        for (int i = 0; i < arrays.length; i++) {
            if (arrays[i].length > 0) {
                minHeap.add(new Element(arrays[i][0], i, 0));
                totalSize += arrays[i].length;
            }
        }
        
        int[] result = new int[totalSize];
        int index = 0;

        // Process the heap
        while (!minHeap.isEmpty()) {
            Element current = minHeap.poll();
            result[index++] = current.value;
            
            // If there's a next element in the same array, add it to the heap
            if (current.elementIndex + 1 < arrays[current.arrayIndex].length) {
                int nextElement = arrays[current.arrayIndex][current.elementIndex + 1];
                minHeap.add(new Element(nextElement, current.arrayIndex, current.elementIndex + 1));
            }
        }
        
        return result;
    }

    private static class Element {
        int value;
        int arrayIndex;
        int elementIndex;

        Element(int value, int arrayIndex, int elementIndex) {
            this.value = value;
            this.arrayIndex = arrayIndex;
            this.elementIndex = elementIndex;
        }
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

class Solution:
    def mergeKSortedArrays(self, arrays: List[List[int]]) -> List[int]:
        min_heap = []
        total_size = 0

        # Initialize the heap with the first element of each array
        for i, array in enumerate(arrays):
            if array:
                heapq.heappush(min_heap, (array[0], i, 0))
                total_size += len(array)

        result = []
        result.extend([0] * total_size)  # Pre-allocate space for the result
        index = 0

        # Process the heap
        while min_heap:
            value, array_index, element_index = heapq.heappop(min_heap)
            result[index] = value
            index += 1

            # If there's a next element in the same array, add it to the heap
            if element_index + 1 < len(arrays[array_index]):
                next_element = arrays[array_index][element_index + 1]
                heapq.heappush(min_heap, (next_element, array_index, element_index + 1))

        return result