Update kth-smallest-element-in-a-sorted-matrix.cpp

kamyu104 · web-flow · commit 3cb7070d7cab · 2016-08-01T22:37:35.000+08:00
diff --git a/C++/kth-smallest-element-in-a-sorted-matrix.cpp b/C++/kth-smallest-element-in-a-sorted-matrix.cpp
@@ -1,34 +1,36 @@
-// Time:  O(klogk)
-// Space: O(k)
+// Time:  O(k * log(min(n, m, k))), with n x m matrix
+// Space: O(min(n, m, k))
 
 class Solution {
 public:
     int kthSmallest(vector<vector<int>>& matrix, int k) {
         int kth_smallest = 0;
 
-        using P = pair<int, int>;
-        const auto Compare = [&matrix](const P& a, const P& b) {
-            return matrix[a.first][a.second] > matrix[b.first][b.second];
-        };
-
-        priority_queue<P, vector<P>, decltype(Compare)> min_heap(Compare);
-        min_heap.emplace(0, 0);
-
-        for (int i = 0; i < k; ++i) {
-            const auto idx = min_heap.top();
-            min_heap.pop();
-            
-            if (idx.first == 0 && idx.second + 1 < matrix[0].size()) {
-                min_heap.emplace(0, idx.second + 1);
+        using P = pair<int, pair<int, int>>;
+        priority_queue<P, vector<P>, greater<P>> q;
+        auto push = [&matrix, &q](int i, int j) {
+            if (matrix.size() > matrix[0].size()) {
+                if (i < matrix[0].size() && j < matrix.size()) {
+                    q.emplace(matrix[j][i], make_pair(j, i));
+                }
+            } else {
+                if (i < matrix.size() && j < matrix[0].size()) {
+                    q.emplace(matrix[i][j], make_pair(i, j));
+                }
             }
+        };
 
-            if (idx.first + 1 < matrix.size()) {
-                min_heap.emplace(idx.first + 1, idx.second);
+        push(0, 0);
+        while (!q.empty() && k--) {
+            auto tmp = q.top(); q.pop();
+            kth_smallest = tmp.first;
+            int i, j;
+            tie(i, j) = tmp.second;
+            push(i, j + 1);
+            if (j == 0) {
+                push(i + 1, 0);
             }
-
-            kth_smallest = matrix[idx.first][idx.second];
         }
-
-        return kth_smallest;
+        return kth_smallest; 
     }
 };
