Description

Submission

class Solution {
public:
    vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
        int n = mat.size();
        for(int i = 0; i < n; ++i) {
            mat[i].push_back(i);
        }

        sort(mat.begin(), mat.end(), [](vector<int>& r1, vector<int>& r2) {
            int cnt1 = upper_bound(r1.begin(), r1.end() - 1, 1, greater<int>()) - r1.begin();
            int cnt2 = upper_bound(r2.begin(), r2.end() - 1, 1, greater<int>()) - r2.begin();
            if(cnt1 == cnt2) return r1.back() < r2.back();
            return cnt1 < cnt2;
        });

        vector<int> rets;

        for(int i = 0; i < k; ++i) {
            rets.push_back(mat[i].back());
        }

        return rets;
    }
};

class Solution {
public:
    vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
        int n = mat.size();
        vector<pair<int,int>> sums(n);

        for(int i = 0; i < mat.size(); ++i) {
            sums[i] = {accumulate(mat[i].begin(), mat[i].end(), 0), i};
        }

        sort(sums.begin(), sums.end());

        vector<int> rets;
        for(int i = 0; i < k; ++i) {
            rets.push_back(sums[i].second);
        }

        return rets;
    }
};