1369C - RationalLee

 Problem Link

Algorithm Used: greedy

TLDR: sort a[] in ascending order and then allocate the largest k numbers among the groups, ensuring the the groups with w[i] = 1 get the largest of these. Then iterate over the groups in descending order of w[i] and allocate the smallest (w[i] - 1) numbers to the ith group. Add the maximum and minimum numbers allocated for every group to the answer.

Long answer:

In this problem, we need to split $N\; integers\; among\; M\; group\; such\; that\; the\; ith\; group\; contains\; W[i]\; numbers.\; We\; need\; to\; find\; the\; maximum\; sum\; of\; values\; possible\; by\; splitting\; the\; numbers\; optimally,\; where\; the\; value\; of\; a\; group\; is\; the\; sum\; of\; the\; maximum\; and\; minimum\; numbers\; in\; it.$

Intuitively, we wouldn't want to "waste" the larger numbers by putting them in a group where they don't contribute to the answer. So to prevent this, we can split the larger numbers across the groups. Since there can only be a maximum of k groups, we will split the largest k numbers among these groups. Particularly, we would put the largest numbers in the groups that require only one number, since the value of these groups is 2 times the value of the singular number put in it.

On the other hand, we would like to "waste" the smaller numbers since these only bring down the answer. We can do this by allocating (w[i] - 1) consecutive numbers at a time to a group so that (w[i] - 2) numbers are "wasted".

One important factor to consider is the order of groups to which we are allocating numbers. We want to add numbers in descending order of the sizes of groups so that we can "waste" more of the smaller numbers. We can see this with an example:

Suppose we have these 5 numbers:

-10 -9 -3 5 6

and we need to add these to two groups of sizes 3 and 2

If we allocated to the group of size 2 first, then the smallest numbers in each of the groups would be -10 and -3, resulting in a smaller value of the second group. However, if we added the numbers to the group of size 3 first, then we would get to "waste" the numbers -9 and -3 and let the smallest numbers in each of the groups be -10 and 5 resulting in a larger value of the group.

Code

#include <iostream>

#include <vector>

#include <set>

#include <algorithm>

using namespace std;

using ll = long long;

int main() {

    // ios_base::sync_with_stdio(false);

    // cin.tie(0);

    // cout.tie(0);

    int T;

    cin >> T;

    while (T--) {

        int n, k;

        cin >> n >> k;

        vector <int> a;

        for (int i = 0; i < n; i++) {

            int x;

            cin >> x;

            a.push_back(x);

        }

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

        vector <int> w(k);

        for (int i = 0; i < k; i++)

            cin >> w[i];

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

        vector <pair <int, int>> allocs(k);

        ll ans = 0;

        for (int i = 0; i < k; i++) {

            ans += a.back();

            allocs[i].first = a.back();

            a.pop_back();

        }      

        for (int i = k-1, p = 0; i >= 0; i--) {

            if (w[i] == 1) {

                ans += allocs[i].first;

                allocs[i].second = allocs[i].second;

                continue;

            } else {

                allocs[i].second = a[p];

                ans += a[p];

                p = p + w[i] - 1;

            }

        }

        cout << ans << '\n';

    }

}