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

时间:2022-03-10 12:18:47

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