USACO 2023 US Open Contest Silver Division - Milk Sum

Problem link: here

Solution Author: Stefan Dascalescu

Problem Solution

Sort the original array and then do prefix sums on it so that you can find the original answer.

For each query, you want to do the following three steps:

• Remove the original value from the array and assume your array has now n-1 values
• Binary search where would the value end up in the new array
• Add the contribution you would get as a result (be careful about the original position to properly discount it from the answer).

Source code

The source code in C++ can be seen below.

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

int main() {

    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int n;
    cin >> n;

    vector <pair<int, int> > v(n+1);
    vector <int> pos(n+1);
    for (int i = 1; i <= n; i++) {
        cin >> v[i].first;
        v[i].second = i;
    }

    sort(v.begin() + 1, v.begin() + n + 1);
    for (int i = 1; i <= n; i++) {
        pos[v[i].second] = i;
    }

    vector <long long> prefix_sum (n+1);
    for (int i = 1; i <= n; i++) {
        prefix_sum[i] = prefix_sum[i-1] + v[i].first;
    }

    long long sum = 0;
    for (int i = 1; i <= n; i++) {
        sum += 1LL * i * v[i].first;
    }

    int q;
    cin >> q;

    for (int i = 1; i <= q; i++) {
        int tpos, val;
        cin >> tpos >> val;


        long long ans = sum - 1LL * pos[tpos] * v[pos[tpos]].first; // subtract contrib of current value        
        ans -= prefix_sum[n] - prefix_sum[pos[tpos]]; // subtract prefix sum of positions to the right

        int L = 1;
        int R = n;
        int target_pos = 0;

        while (L <= R) {
            int mid = (L + R) / 2;
            if (v[mid].first <= val) {
                target_pos = mid;
                L = mid + 1;
            }
            else {
                R = mid - 1;
            }
        }

        ans += 1LL * (target_pos + 1 - (target_pos >= pos[tpos])) * val; // adding back contrib of current value
        ans += prefix_sum[n] - prefix_sum[target_pos]; // adding back new contrib
        if (pos[tpos] > target_pos) { // remove the original value if needed
            ans -= v[pos[tpos]].first;
        }
        cout << ans << '\n';
    }    
    return 0;
}