557C – Arthur and Table (Codeforces)

by afruizc

The problem statement can be found here.

Problem:
Given the length and the energy to remove a set of legs from a table, we are to balance this table in such a way that we minimize the amount of energy spent. A table is balanced when the amount of legs of maximum size is greater than half of the total amount of legs. For example, a table with 5 legs of lengths, 3 1 4 4 4, will be balanced if:

  • We remove legs 1 and 3 (we keep legs of size 4). or
  • We remove legs of length 3 and 4 (we keep leg of size 1). or
  • We remove legs of length 1 and 4 (we keep leg of size 3).

Solution:
We will use a simple algorithm to solve this problem: we will first group the legs by length. Then, for each one of these groups, we will assume it is the one that will balance the table. We will then proceed to remove legs outside of the group in such a way that we minimize the energy spent. Finally, over all the groups, we pick the one that yields the lowest energy expenditure.

Now what? Well, we still need to figure out the details of the implementation. To group the legs, we can sort them by length in non-decreasing order. Let’s let g_i be the group that is being currently processed and k the size of such group (the number of legs with the same length). Because we are assuming that g_i is the group that will balance the table, we need to remove all the legs that have length greater than it. This can be done in constant time by precomputing the sum from right to left after having sorted the array. As for the legs with length less than g_i, we need to keep a total of k-1 legs, such that the energy removing the other legs is minimal. Every time we finish dealing with a group of legs, we can add them to a sorted list (multiset), so that it is easier to select which legs should be unscrewed. Let’s call this list s. Now, there is two ways we can select the optimal subset:

Select legs with minimal energy to unscrew: (TLE response)
We can traverse s from left to right (assuming it is sorted in non-decreasing order), and add up the first min(0, s.length - (k-1)) elements. The problem with this approach is that we can go up to quadratic complexity, because we have iterate over all the legs that are to be removed. As an example, imagine a list of leg lengths that are all different, i.e. each group will only have one element. For each group (each leg in this case), we will then have to iterate over s.length legs and add the energies up. After a group (a leg) is dealt with, we will add this leg to s. This clearly suggests that we will approach quadratic complexity (which TLEs). Which leads us to the next approach…

Select legs with maximal energy and unscrew the rest: (AC response)
We can use a similar idea, and, helped by keeping a sum of the elements added to s, we can instead select the legs with maximal energy. To this end, we will have s be a non-increasing sorted array. This approach improves the complexity because we will only have to iterate through at most k-1 elements of s every time. Clearly, this approach is not quadratic. In this case, the total energy spent by removing legs to the left of g_i will be equal to: (sum(s) - sum(s[:(k-1)])) (I’m using slicing to simplify the notation).

This then finishes our analysis. If there are any questions or concerns, just leave a comment.

Code:

#include <cstdio>
#include <set>
#include <algorithm>
#include <functional>

#define Max(a, b) ((a) > (b) ? (a) : (b))
#define Min(a, b) ((a) < (b) ? (a) : (b))

#define FOR(i, s, e, inc) for (int i=(s) ; i<(int)(e) ; i+=(inc))
#define RFOR(i, s, e, dec) for (int i=(s) ; i>=(int)(e) ; i-=(dec))
#define REP(i, e) FOR(i, 0, e, 1)
#define RREP(i, s) RFOR(i, s, 0, 1)

using namespace std;

int n;
int sum[500001];
pair<int, int> legs[500000];
multiset<int, greater<int>> past_legs;
int energy_so_far = 0;

void find_left_sum() {
    sum[n] = 0;
    sum[n-1] = legs[n-1].second;
    RREP(i, n-2) {
        sum[i] = sum[i + 1] + legs[i].second;
    }
}

int solve_one(int k, int init, int end) {
    int total_energy = sum[end + 1] + energy_so_far;
    auto it = past_legs.begin();
    for (int i=0 ; i<(k-1) && it != past_legs.end() ; i++) {
        total_energy -= (*it);
        it++;
    }

    for (int i=init ; i<=end ; i++) {
        past_legs.insert(legs[i].second);
        energy_so_far += legs[i].second;
    }

    return total_energy;
}

int main() {
    while (scanf("%d\n", &n) == 1) {
        past_legs.clear();
        energy_so_far = 0;
        REP(i, n) {
            scanf("%d", &legs[i].first);
        }

        REP(i, n) {
            scanf("%d", &legs[i].second);
        }

        sort(legs, legs + n);
        find_left_sum();

        int k, init, cur_l;
        int ans = 1e9;
        k = 1;
        cur_l = legs[0].first;
        init = 0;
        for (int i=1 ; i<n ; i++) {
            if (legs[i].first != cur_l) {
                // Process the set of legs
                int energy = solve_one(k, init, i-1);
                ans = Min(ans, energy);
                k = 1;
                init = i;
                cur_l = legs[i].first;
            }
            else {
                k++;
            }
        }

        int energy = solve_one(k, init, n-1);
        printf("%d\n", Min(ans, energy));
    }
}