Merge k Sorted Arrays【合并k个有序数组】【优先队列】

Given k sorted integer arrays, merge them into one sorted array.
Example Given 3 sorted arrays:
[
[1, 3, 5, 7],
[2, 4, 6],
[0, 8, 9, 10, 11]
]
return [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11].
Challenge Do it in O(N log k).
N is the total number of integers. k is the number of arrays.

---

Solution:
创建一个大小为N的数组保存最后的结果
数组本身已经从小到大排好序，所以我们只需创建一个大小为k的最小堆，堆中初始元素为k个数组中的每个数组的第一个元素，每次从堆中取出最小元素（堆顶元素），并将其存入输出数组中，将堆顶元素所在行的下一元素加入堆，重新排列出堆顶元素，时间复杂度为logk，总共N个元素，所以总体时间复杂度是Nlogk

import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.Queue;

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 Main {
	
	private Comparator<Element> MyCmp = new Comparator<Element>() {

		@Override //升序
		public int compare(Element o1, Element o2) {
			return o1.val - o2.val;
		}
	};
	
	
	public int[] mergeKSortedArrays(int[][] arr) {
		if(arr == null) {
			return new int[0];
		}
		
		int sum = 0;
		Queue<Element> q = new PriorityQueue<Element>(
				arr.length,MyCmp);
		
		for(int i = 0; i < arr.length; i++) {
			if(arr[i].length > 0) {
				Element e = new Element(i, 0, arr[i][0]);
				q.add(e);
				sum += arr[i].length; //记录结果数组总长度
			}
		}
		
		int[] res = new int[sum];
		int idx = 0;
		while(!q.isEmpty()) {
			Element e = q.poll(); //弹出堆顶最小值
			res[idx++] = e.val;
                        // 当前结点被 PriorityQueue 抛出来后，当前行的第二个结点加入 PriorityQueue
			if(e.col + 1 < arr[e.row].length) {     //将当前最小所在数组的下一个元素存入pq
				q.add(new Element(e.row, e.col+1, arr[e.row][e.col+1]));
			}
		}
		return res;
	}
	public static void main(String[] args) {
		Main m = new Main();
		int[][] arr = {{1, 3, 5, 7},{2, 4, 6},{0, 8, 9, 10, 11}};
		int[] res = m.mergeKSortedArrays(arr);
		for(int i=0; i<res.length; i++) {
			System.out.print(res[i] + " ");
		}
	}
}
/*
 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]
 */

GG手写堆排序的写法；

// Java program to merge k sorted 
// arrays of size n each. 

// A min heap node 
class MinHeapNode 
{ 
	int element; // The element to be stored 
	
	// index of the array from 
	// which the element is taken 
	int i; 
	
	// index of the next element 
	// to be picked from array 
	int j; 

	public MinHeapNode(int element, int i, int j) 
	{ 
		this.element = element; 
		this.i = i; 
		this.j = j; 
	} 
}; 

// A class for Min Heap 
class MinHeap 
{ 
	MinHeapNode[] harr; // Array of elements in heap 
	int heap_size; // Current number of elements in min heap 

	// Constructor: Builds a heap from 
	// a given array a[] of given size 
	public MinHeap(MinHeapNode a[], int size) 
	{ 
		heap_size = size; 
		harr = a; 
		int i = (heap_size - 1)/2; 
		while (i >= 0) 
		{ 
			MinHeapify(i); 
			i--; 
		} 
	} 

	// A recursive method to heapify a subtree 
	// with the root at given index This method 
	// assumes that the subtrees are already heapified 
	void MinHeapify(int i) 
	{ 
		int l = left(i); 
		int r = right(i); 
		int smallest = i; 
		if (l < heap_size && harr[l].element < harr[i].element) 
			smallest = l; 
		if (r < heap_size && harr[r].element < harr[smallest].element) 
			smallest = r; 
		if (smallest != i) 
		{ 
			swap(harr, i, smallest); 
			MinHeapify(smallest); 
		} 
	} 

	// to get index of left child of node at index i 
	int left(int i) { return (2*i + 1); } 

	// to get index of right child of node at index i 
	int right(int i) { return (2*i + 2); } 

	// to get the root 
	MinHeapNode getMin() 
	{ 
		if(heap_size <= 0) 
		{ 
			System.out.println("Heap underflow"); 
			return null; 
		} 
		return harr[0]; 
	} 

	// to replace root with new node 
	// "root" and heapify() new root 
	void replaceMin(MinHeapNode root) { 
		harr[0] = root; 
		MinHeapify(0); 
	} 

	// A utility function to swap two min heap nodes 
	void swap(MinHeapNode[] arr, int i, int j) { 
		MinHeapNode temp = arr[i]; 
		arr[i] = arr[j]; 
		arr[j] = temp; 
	} 

	// A utility function to print array elements 
	static void printArray(int[] arr) { 
		for(int i : arr) 
			System.out.print(i + " "); 
		System.out.println(); 
	} 

	// This function takes an array of 
	// arrays as an argument and All 
	// arrays are assumed to be sorted. 
	// It merges them together and 
	// prints the final sorted output. 
	static void mergeKSortedArrays(int[][] arr, int k) 
	{ 
		MinHeapNode[] hArr = new MinHeapNode[k]; 
		int resultSize = 0; 
		for(int i = 0; i < arr.length; i++) 
		{ 
			MinHeapNode node = new MinHeapNode(arr[i][0],i,1); 
			hArr[i] = node; 
			resultSize += arr[i].length; 
		} 

		// Create a min heap with k heap nodes. Every heap node 
		// has first element of an array 
		MinHeap mh = new MinHeap(hArr, k); 

		int[] result = new int[resultSize]; // To store output array 

		// Now one by one get the minimum element from min 
		// heap and replace it with next element of its array 
		for(int i = 0; i < resultSize; i++) 
		{ 

			// Get the minimum element and store it in result 
			MinHeapNode root = mh.getMin(); 
			result[i] = root.element; 

			// Find the next element that will replace current 
			// root of heap. The next element belongs to same 
			// array as the current root. 
			if(root.j < arr[root.i].length) 
				root.element = arr[root.i][root.j++]; 
			// If root was the last element of its array 
			else
				root.element = Integer.MAX_VALUE; 

			// Replace root with next element of array 
			mh.replaceMin(root); 
		} 

		printArray(result); 

	} 

	// Driver code 
	public static void main(String args[]){ 
		int[][] arr= {{2, 6, 12, 34}, 
				{1, 9, 20, 1000}, 
				{23, 34, 90, 2000}}; 

		System.out.println("Merged array is :"); 

		mergeKSortedArrays(arr,arr.length); 
	} 
}; 

// This code is contributed by shubham96301 

PS：一面头条笔试