Good grades

input/code.cpp: In function 'void solve()':
input/code.cpp:26:23: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   26 |     for (int i = 0; i < nums.size(); i++)
      |                     ~~^~~~~~~~~~~~~
input/code.cpp:45:23: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   45 |     for (int j = 1; j < dp[0].size(); j++)
      |                     ~~^~~~~~~~~~~~~~
input/code.cpp:47:27: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::vector<std::pair<long long int, long long int> > >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   47 |         for (int i = 1; i < dp.size(); i++)
      |                         ~~^~~~~~...

#include <iostream>
#include <vector>
#include <array>
#include <string>
#include <algorithm>
#include <numeric>
#include <unordered_map>
#include <unordered_set>
#include <queue>
#include <climits>
#include <cmath>
#include <functional>
#include <type_traits>
#include <fstream>
#include <bitset>
 
#define int long long
using namespace std;
 

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

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

    /*
    for (int i = 0; i < nums.size(); i++)
        cout << nums[i].first << " ";
    cout << endl;
    //*/

    vector<vector<pair<int, int>>> dp(nums.size()+1, vector<pair<int, int>>(k+1, {LLONG_MIN / 2, 0}));
    vector<vector<int>> dpMem(nums.size()+1, vector<int>(k+1, 0));
    dp[0][0] = {0, 0};

    for (int j = 1; j < dp[0].size(); j++)
    {
        for (int i = 1; i < dp.size(); i++)
        {
            // p1: create new, p2: continue other
            int maxIdx = 0;
            int maxVal = LLONG_MIN / 2;
            for (int t = 0; t < i; t++)
            {
                int cand = dp[t][j-1].first + (i-t) * nums[t].first;
                if (cand > maxVal)
                {
                    maxVal = cand;
                    maxIdx = t;
                }
            }
            dp[i][j] = {maxVal, maxIdx};
            /*
            int p1 = dp[i-1][j-1].first + nums[i-1].first;
            int p2 = dp[i-1][j].first + dp[i-1][j].second;
            if (p1 >= p2)
            {
                dp[i][j] = {p1, nums[i-1].first};
                dpMem[i][j] = 1;
            }
            else
            {
                dp[i][j] = {p2, dp[i-1][j].second};
            }
            */
            //cout << dp[i][j].first << "-" << dp[i][j].second << "-" << dpMem[i][j] << "       ";
        }
        //cout << endl;
    }

    vector<int> res(n);

    int i = n;
    for (int j = k; j > 0; j--)
    {
        int am = dp[i][j].second;
        while (i > am)
        {
            res[nums[i-1].second] = j;
            i--;
        }
    }

    for (int i = 0; i < res.size(); i++)
        cout << res[i] << " ";

}
 
 
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
 
    int t = 1;
 
    //cin >> t;
    //ifstream f("testcase.txt");
    //cin.rdbuf(f.rdbuf());
 
    for (int i = 0; i < t; i++)
        solve();
 
    return 0;
}

Test 1

Verdict: ACCEPTED

Test 2

Verdict: ACCEPTED

Test 3

Verdict: ACCEPTED

Test 4

Verdict: ACCEPTED

Test 5

Verdict: ACCEPTED

Test 6

Verdict: ACCEPTED

Test 7

Verdict: ACCEPTED

Test 8

Verdict: ACCEPTED

Test 9

Verdict: ACCEPTED

Test 10

Verdict: ACCEPTED

Test 11

Verdict: ACCEPTED

Test 12

Verdict: ACCEPTED

Test 13

Verdict: ACCEPTED

Test 14

Verdict: ACCEPTED

Test 15

Verdict: ACCEPTED

Test 16

Verdict: ACCEPTED

Test 17

Verdict: ACCEPTED

Test 18

Verdict: ACCEPTED

Test 19

Verdict: ACCEPTED

Test 20

Verdict: ACCEPTED

Test 21

Verdict: ACCEPTED

Test 22

Verdict: ACCEPTED

Test 23

Verdict: ACCEPTED

Test 24

Verdict: ACCEPTED

Test 25

Verdict: ACCEPTED

Test 26

Verdict: ACCEPTED

Test 27

Verdict: ACCEPTED

Test 28

Verdict: ACCEPTED

Test 29

Verdict: ACCEPTED

Test 30

Verdict: ACCEPTED

Test 31

Verdict: ACCEPTED

Test 32

Verdict: ACCEPTED

Test 33

Verdict: ACCEPTED

Test 34

Verdict: ACCEPTED

Test 35

Verdict: ACCEPTED

Test 36

Verdict: ACCEPTED

Test 37

Verdict: ACCEPTED

Test 38

Verdict: ACCEPTED

Test 39

Verdict: ACCEPTED

Test 40

Verdict: ACCEPTED

Test 41

Verdict: ACCEPTED

Test 42

Verdict: ACCEPTED

Test 43

Verdict: ACCEPTED

Test 44

Verdict: ACCEPTED

Test 45

Verdict: ACCEPTED

Test 46

Verdict: ACCEPTED

Test 47

Verdict: ACCEPTED

Test 48

Verdict: ACCEPTED

Test 49

Verdict: ACCEPTED

Test 50

Verdict: ACCEPTED

Test 51

Verdict: ACCEPTED

Test 52

Verdict: ACCEPTED

Test 53

Verdict: ACCEPTED

Test 54

Verdict: ACCEPTED

Test 55

Verdict: ACCEPTED

Test 56

Verdict: ACCEPTED

Test 57

Verdict: ACCEPTED

Test 58

Verdict: ACCEPTED

Test 59

Verdict: ACCEPTED