Update the-k-weakest-rows-in-a-matrix.cpp

Author: kamyu104

C++/the-k-weakest-rows-in-a-matrix.cpp

@@ -13,7 +13,7 @@ class Solution {
                 }
                 lookup.emplace(i);
                 result.emplace_back(i);
-                if (!--k) {
+                if (result.size() == k) {
                     return result;
                 }
             }
@@ -24,7 +24,7 @@ class Solution {
             }
             lookup.emplace(i);
             result.emplace_back(i);
-            if (!--k) {
+            if (result.size() == k) {
                 break;
             }
         }
@@ -44,7 +44,7 @@ class Solution2 {
                     continue;
                 }
                 lookup[i] = true;
-                if (!--k) {
+                if (lookup.size() == k) {
                     return lookup.keys();
                 }
             }
@@ -54,7 +54,7 @@ class Solution2 {
                 continue;
             }
             lookup[i] = true;
-            if (!--k) {
+            if (lookup.size() == k) {
                 break;
             }
         }
@@ -91,6 +91,10 @@ class Solution2 {
                       });
             return result;
         }
+        
+        int size() const {
+            return list_.size();
+        }
 
     private:
         list<pair<K, V>> list_;