The key ideas:

  • Use a std::pair to hold {count, index}, so it can compare count first then the index.
  • Use a min heap priority queue to get the K weakest rows.
class Solution {
public:
	vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
		// min heap
		priority_queue<
			std::pair<int, int>,
			vector<std::pair<int, int> >,
			std::greater<std::pair<int, int> > > pq;
		for (auto iter = mat.begin(); iter != mat.end(); ++iter) {
			int c = count((*iter).begin(), (*iter).end(), 1);
			pq.push({c, iter - mat.begin()});
		}
		vector<int> ret;
		for (; k > 0; --k) {
			ret.push_back(pq.top().second);
			pq.pop();
		}
		return ret;
	}
};