Code : Learning to Paint

#include <bits/stdc++.h>
using namespace std;

#define endl "\n"

int k;

void add(multiset<long long> &mset, long long val) {
    mset.insert(val);
    if ((int)mset.size() > k) {
        mset.erase(mset.begin());
    }
}

void solve() {
    int n;
    cin >> n >> k;
    vector<vector<int>> a(n, vector<int>(n, 0));
    for (int i = 0; i < n; i++) {
        for (int j = i; j < n; j++) {
            cin >> a[i][j];
        }
    }

    vector<multiset<long long>> indp(n), outdp(n);
    add(indp[0], a[0][0]), add(outdp[0], 0);
    for (int i = 1; i < n; i++) {
        for (auto &ele : indp[i - 1])
            add(outdp[i], ele);
        for (auto &ele : outdp[i - 1])
            add(outdp[i], ele);

        // Look for last white vertex.
        add(indp[i], a[0][i]);
        for (int j = i - 1; j >= 0; j--)
            for (auto &ele : outdp[j])
                add(indp[i], ele + a[j + 1][i]);
    }

    multiset<long long> mans;
    for (auto &ele : indp.back())
        add(mans, ele);

    for (auto &ele : outdp.back())
        add(mans, ele);

    for (auto itr = mans.rbegin(); itr != mans.rend(); itr++)
        cout << *itr << " ";
    cout << endl;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int t;
    cin >> t;
    for (int zz = 0; zz < t; zz++) {
        solve();
    }
    return 0;
}

#include <bits/stdc++.h>
using namespace std;

#define endl "\n"

using st = pair<multiset<long long>::reverse_iterator, pair<long long, int>>;

int k;

void add(multiset<long long> &mset, long long val) {
    mset.insert(val);
    if ((int)mset.size() > k) {
        mset.erase(mset.begin());
    }
}

void solve() {
    int n;
    cin >> n >> k;
    vector<vector<int>> a(n, vector<int>(n, 0));
    for (int i = 0; i < n; i++) {
        for (int j = i; j < n; j++) {
            cin >> a[i][j];
        }
    }

    vector<multiset<long long>> indp(n), outdp(n);
    add(indp[0], a[0][0]);
    add(outdp[0], 0);
    for (int i = 1; i < n; i++) {
        for (auto &ele : indp[i - 1])
            add(outdp[i], ele);
        for (auto &ele : outdp[i - 1])
            add(outdp[i], ele);

        // {mitr, delta, j}
        auto func = [](st small, st big) {
            auto sitr = small.first;
            auto bitr = big.first;

            auto sdelta = small.second.first;
            auto bdelta = big.second.first;

            return *sitr + sdelta > *bitr + bdelta;
        };

        multiset<st, decltype(func)> choices(func);
        for (int j = i - 1; j >= 0; j--) {
            auto now = st{outdp[j].rbegin(), {a[j + 1][i], j}};
            choices.insert(now);
        }

        while (!choices.empty() && indp[i].size() < k) {
            auto now = *choices.begin();
            choices.erase(choices.begin());
            auto itr = now.first;
            auto delta = now.second.first;
            int j = now.second.second;

            add(indp[i], *itr + delta);
            now.first++;
            if (now.first != (outdp[j].rend())) {
                choices.insert(now);
            }
        }

        // Look for last white vertex.
        add(indp[i], a[0][i]);
    }

    multiset<long long> mans;
    for (auto &ele : indp.back())
        add(mans, ele);
    for (auto &ele : outdp.back())
        add(mans, ele);

    for (auto itr = mans.rbegin(); itr != mans.rend(); itr++) {
        cout << *itr << " ";
    }
    cout << endl;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int t;
    cin >> t;
    for (int zz = 0; zz < t; zz++) {
        solve();
    }
    return 0;
}