LintCode 486 Merge k Sorted Arrays

思路

要求时间复杂度O(N logk)
使用priorityqueue组成最小堆。堆的大小为k，首先将所有数组的第一个元素放进去，需要记录下是哪一个数组的第几个元素。然后每次从pq中拿出最小的元素加入答案序列，然后根据拿出来的那个元素的row和col，将与其同数组的下一个元素加入pq，保证pq的大小始终是k。
【实现：可以新建一个class去存储每一个元素所在的row，col和val，然后对于pq重写一个比较器即可】

复杂度

时间复杂度O(N logk)，空间复杂度O(N)，N为所有元素的个数

代码

class Element {
    public int row, col, val;
    public Element(int row, int col, int val) {
        this.row = row;
        this.col = col;
        this.val = val;
    }
}

public class Solution {
    /**
     * @param arrays: k sorted integer arrays
     * @return: a sorted array
     */
    public int[] mergekSortedArrays(int[][] arrays) {
        // write your code here
        // min-heap O(Nlogk)
        Comparator<Element> cmp = new Comparator<Element>() {
            @Override
            public int compare(Element e1, Element e2) {
                return e1.val - e2.val;
            }
        };
        
        int k = arrays.length;
        PriorityQueue<Element> pq = new PriorityQueue<>(k, cmp);
        
        int size = 0;
        for(int i = 0; i < k; i++) {
            if(arrays[i].length > 0) {
                Element e = new Element(i, 0, arrays[i][0]);
                pq.offer(e);
                size += arrays[i].length;
            }
        }
        
        int[] ans = new int[size];
        int index = 0;
        while(!pq.isEmpty()) {
            Element tmp = pq.poll();
            ans[index++] = tmp.val;
            
            int next_row = tmp.row, next_col = tmp.col + 1;
            if(next_row < arrays.length && next_col < arrays[next_row].length) {
                Element next = new Element(next_row, next_col, arrays[next_row][next_col]);
                pq.offer(next);
            }
        }
        return ans;
    }
}