378. 有序矩阵中第K小的元素

/**
 * 378. 有序矩阵中第K小的元素
 * @author wsq
 * @date 2020/10/09
	给定一个 n x n 矩阵，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。
	请注意，它是排序后的第 k小元素，而不是第 k 个不同的元素。
	
	示例：
	matrix = [
	   [ 1,  5,  9],
	   [10, 11, 13],
	   [12, 13, 15]
	],
	k = 8,
	返回 13。
	
	链接：https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrix
 */
package search;

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

public class KthSmallest {
	/**
	 *	找到有序队列中第第K小的元素，包含重复的
	 *	采用归并排序的思想，找到第K个最小值结束
	 *	由于是n个有序队列归并排序，为了优化检索速度，采用小根堆维护，来优化时间复杂度
	 * @param matrix
	 * @param k
	 * @return
	 */
	public int kthSmallest(int[][] matrix, int k) {
		PriorityQueue<int[]> pq = new PriorityQueue<>(new Comparator<int[]>() {
			public int compare(int[] o1, int[] o2) {
				return o1[0] - o2[0];
			};
		});
		
		int n = matrix.length;
		
		for(int i = 0; i < n; ++i) {
			pq.offer(new int[] {matrix[i][0], i, 0});
		}
		
		for(int i = 0; i < k - 1; i++) {
			int[] now = pq.poll();
			if(now[2] != n - 1) {
				pq.offer(new int[] {matrix[now[1]][now[2] + 1], now[1], now[2] + 1});
			}
		}
		return pq.poll()[0];
    }
	
	public static void main(String[] args) {
		int[][] matrix = {{1,  5,  9},
		                  {10, 11, 13},
		                  {12, 13, 15}};
		int k = 8;
		KthSmallest ks = new KthSmallest();
		int ans = ks.kthSmallest(matrix, k);
		System.out.println(ans);
	}
}