Merge k Sorted Arrays

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

Thoughts

堆排序，用一个自建数据结构以记录下一个要加入堆的位置。

Code

public class Solution {
    /**
     * @param arrays k sorted integer arrays
     * @return a sorted array
     */
    public List<Integer> mergekSortedArrays(int[][] arrays) {
        PriorityQueue<Element> pq = new PriorityQueue<>(1, new ElementComparator());
        int k = arrays.length;
        List<Integer> res = new ArrayList<>();
        for (int i = 0; i < k; i++) {
            if (arrays[i].length > 0) {
                pq.add(new Element(arrays[i][0], i, 0));
            }
        }

        while (pq.size() != 0) {
            Element min = pq.poll();
            res.add(min.val);
            if (min.y + 1 < arrays[min.x].length) {
                pq.add(new Element(arrays[min.x][min.y + 1], min.x, min.y + 1));
            }
        }

        return res;
    }

    class Element {
        int val;
        int x, y;
        Element(int val, int x, int y) {
            this.val = val;
            this.x = x;
            this.y = y;
        }
    }

    class ElementComparator implements Comparator<Element> {
        public int compare(Element a, Element b) {
            return a.val - b.val;
        }
    }


}

Analysis

TC: O(nlgk), 堆里维持k个元素，一共添加和删除n次，每次插入删除复杂度为O(lgk).