#include "header.hpp"
class Solution {
public:
std::vector<int> kWeakestRows(std::vector<std::vector<int>> &mat, int k) {
std::pair<std::vector<int>, std::vector<int>> sums;
sums.first.reserve(mat.size());
sums.second.reserve(mat.size());
int row = 0;
int sum = 0;
for (auto &iter : mat) {
for (auto &ya_iter : iter) {
sum += ya_iter;
if (0 == ya_iter) {
break;
}
}
sums.first[row] = row;
sums.second[row] = sum;
row++;
sum = 0;
}
for (int i = 0; i < row - 1; ++i) {
for (int j = 0; j < row - 1; ++j) {
if (sums.second[j] > sums.second[j + 1]) {
std::swap(sums.first[j], sums.first[j + 1]);
std::swap(sums.second[j], sums.second[j + 1]);
}
}
}
/* for (int i = 0; i < row; ++i) { */
/* std::cout << sums.first[i] << ":" << sums.second[i] << "\n"; */
/* } */
std::vector<int> result(k);
std::copy(sums.first.begin(), sums.first.begin() + k, result.begin());
return result;
}
};
int main(int argc, char **argv) {
Solution solution;
std::vector<std::vector<int>> mat = {{1, 1, 0, 0, 0},
{1, 1, 1, 1, 0},
{1, 0, 0, 0, 0},
{1, 1, 0, 0, 0},
{1, 1, 1, 1, 1}};
auto result = solution.kWeakestRows(mat, 3);
for (int i = 0; i < 3; ++i) {
std::cout << result[i] << "\n";
}
return 0;
}