[LeetCode] 1337. The K Weakest Rows in a Matrix

https://leetcode.com/problems/the-k-weakest-rows-in-a-matrix/

The K Weakest Rows in a Matrix - LeetCode

The K Weakest Rows in a Matrix

0과 1로만 주어진 NxN 행렬이 주어져있다.

1은 무조건 0의 왼쪽에 있어야한다.

이때 1은 각 행의 병사를 나타내는데,

병사가 적은 행 (약한행)부터 K번째 행까지 출력해주면 된다.

이때 병사수가 같으면 행 번호가 작은게 약한 행이다.

 

풀이는 간단하다!

각 병사수로 (배열, 행번호) pair를 생성해주고 sort 해주면 된다.

pair로 sort 시에 compare function을 만드는 방법이 필요하다.

 

전체 코드는 아래와 같다.

 

class Solution {
public:
    vector<int> kWeakestRows(vector<vector<int>>& mat, int k) {
        vector<int> ret;
        vector<pair<int,int>> r;
        for (int i=0; i<mat.size(); i++) {
            int sum = 0;
            for (int j=0; j<mat[i].size() && mat[i][j]==1; j++) sum++;
            r.push_back(make_pair(sum,i));
        }
        sort(r.begin(), r.end(), compare);
        for (int i=0; i<k; i++) ret.push_back(r[i].second);
        return ret;
    }
    static bool compare(pair<int, int> a, pair<int, int> b) {return a.first==b.first?a.second<=b.second:a.first<b.first;}
};