486. 合并k个排序数组

描述

将k个排序数组合并为一个大的排序数组。

样例

给出下面的3个排序数组：

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

合并后的大数组应为：

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

挑战

在 O(Nlogk) 的时间复杂度内完成：

• N是所有数组包含的整数个数。
• k 是数组的个数。

思路

1. 最简单的方法是创建一个N大小的数组，然后把所有数字拷贝进去，然后再进行时间复杂度为O(NlogN)排序算法，这样总体时间复杂度为O(NlogN)
2. 可以利用最小堆完成，时间复杂度是O(Nlogk)，具体过程如下：
   创建一个大小为N的数组保存最后的结果
   数组本身已经从小到大排好序，所以我们只需创建一个大小为k的最小堆，堆中初始元素为k个数组中的每个数组的第一个元素，每次从堆中取出最小元素（堆顶元素），并将其存入输出数组中，将堆顶元素所在行的下一元素加入堆，重新排列出堆顶元素，时间复杂度为logk，总共N个元素，所以总体时间复杂度是Nlogk

代码

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

public class Solution {
    // 从小到大排序
    private Comparator ElementComparator = new Comparator() {
        public int compare(Element left, Element right) {
            return left.val - right.val;
        }
    };
    
    /**
     * @param arrays k sorted integer arrays
     * @return a sorted array
     */
    public int[] mergekSortedArrays(int[][] arrays) {
        if (arrays == null) {
            return new int[0];
        }
        
        int total_size = 0;
        // 默认由小到大顺序
        Queue Q = new PriorityQueue(
            arrays.length, ElementComparator);
        
        // 初始化
        // 把每一行的第一个元素加入 PriorityQueue
        // 顺便统计元素总量
        for (int i = 0; i < arrays.length; i++) {
            // 当前行长度不为 0
            if (arrays[i].length > 0) {
                Element elem = new Element(i, 0, arrays[i][0]);
                Q.add(elem);
                total_size += arrays[i].length;
            }
        }
        
        int[] result = new int[total_size];
        int index = 0;
        while (!Q.isEmpty()) {
            Element elem = Q.poll();
            result[index++] = elem.val;
            // 当前结点被 PriorityQueue 抛出来后，当前行的第二个结点加入 PriorityQueue
            if (elem.col + 1 < arrays[elem.row].length) {
                elem.col += 1;
                elem.val = arrays[elem.row][elem.col];
                Q.add(elem);
            }
        }
        
        return result;
    }
}