Lintcode 486. Merge K Sorted Arrays

Given k sorted integer arrays, merge them into one sorted array.

Example

Example 1:

Input: 
  [
    [1, 3, 5, 7],
    [2, 4, 6],
    [0, 8, 9, 10, 11]
  ]
Output: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]

Example 2:

Input:
  [
    [1,2,3],
    [1,2]
  ]
Output: [1,1,2,2,3]

Challenge

Do it in O(N log k).

• N is the total number of integers.
• k is the number of arrays.

Code (Java):

public class Solution {
    /**
     * @param arrays: k sorted integer arrays
     * @return: a sorted array
     */
    public int[] mergekSortedArrays(int[][] arrays) {
        // write your code here
        if (arrays == null || arrays.length == 0) {
            return new int[0];
        }
        
        int[] ans = mergekSortedArraysHelper(0, arrays.length - 1, arrays);
        
        return ans;
    }
    
    private int[] mergekSortedArraysHelper(int start, int end, int[][] arrays) {
        if (start == end) {
            return arrays[start];
        }
        
        int mid = start + (end - start) / 2;
        int[] left = mergekSortedArraysHelper(start, mid, arrays);
        int[] right = mergekSortedArraysHelper(mid + 1, end, arrays);
        
        return mergeTwoLists(left, right);
    }
    
    private int[] mergeTwoLists(int[] list1, int[] list2) {
        int[] ans = new int[list1.length + list2.length];
        int i = 0;
        int j = 0;
        int k = 0;
        while (i < list1.length && j < list2.length) {
            if (list1[i] < list2[j]) {
                ans[k++] = list1[i++]; 
            } else {
                ans[k++] = list2[j++];
            }
        }
        
        while (i < list1.length) {
            ans[k++] = list1[i++];
        }
        
        while (j < list2.length) {
            ans[k++] = list2[j++];
        }
        
        return ans;
    }
}