378. Kth Smallest Element in a Sorted Matrix

最后更新

二刷

看到这个题有点懵逼。
第一反应是二分，马上觉得不对，又觉得从右上开始（因为做过类似的），也不行。

回头看一刷答案发现居然是，直接PQ。。。。

max Heap就可以了，注意某些时候是可以break的。

到某一个点的时候，能保证比这个点小的元素多于K，就没必要继续添加了。

判断式是(i + 1）* （j + 1) > k

假如K = 7

1 2 3 5
2 3 10 19
3 4 11 20
40 41 44 89

在第二行的时候，只有在19的时候才能保证已经找到的第7小的元素。
在第三行的时候，在11之后才可以保证。
第四行的时候，在41之后就可以保证了。

所以是i+1 j+1，加上这个结果速度也没快多少= =

Time : O((n²)lgK)
Space: O(k)

public class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        PriorityQueue<Integer> pq = new PriorityQueue<>(k, new Comparator<Integer>() {
            public int compare(Integer a, Integer b) {
                return Integer.compare(b, a);
            }
        });
        
        for (int i = 0; i < matrix.length; i++) {
            for (int j = 0; j < matrix[0].length; j++) {
                if ((i+1) * (j+1) > k) break;
                pq.offer(matrix[i][j]);
                if (pq.size() > k) pq.poll();
            }
        }
        return pq.peek();
    }
}

仔细想想，这个题还是一个kth largest(smallest)之类的问题，用PQ做。

有一种做法是二维变一维的MAP方式，也可以，毕竟相邻的坐标是有大小关系的，可以当成一维的find kth largest...不过写起来比较困难。

---

一刷。

做过类似的，一开始的思路被以前模糊的那个题影响了，其实这个题感觉没啥巧办法。

遍历，维护PQ。

需要注意的是1个是PQ里comparator的定义，要反着来，因为我们是维护最小，不是最大，a>b要返还-1

另一个需要注意的是什么时候停止遍历。

以下图为例
1 2 3 5
2 3 10 19
3 4 11 20
40 41 44 89

比如K=7 答案应该是4

4并不是加入PQ的第7个元素

判定是(m+1)*(n+1) > k的时候 往后的元素就不需要再入栈了

所以上图实际入栈的元素是
1 2 3 5
2 3 10
3 4
40

然后看PQ的SIZE，大于K的都拿掉，然后ROOT就是要求的元素。

public int kthSmallest(int[][] matrix, int k) {
        
        PriorityQueue<Integer> pq = new PriorityQueue<Integer>(new Comparator<Integer>(){
        
            public int compare(Integer a, Integer b)
            {
                if(a > b) return -1;
                else if(a < b) return 1;
                else
                    return 0;
            }
        });
        for(int m = 0; m < matrix.length;m++)
        {
            for(int n = 0; n < matrix.length;n++)
            {
                if((m+1)*(n+1) > k)
                {
                    break;
                }
                else
                {
                    pq.add(matrix[m][n]);
                }
            }
            //if(done) break;
        }

        while(pq.size() > k) pq.poll();
        
        return pq.peek();

    }

注意comparator的定义方法