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：一面头条笔试