Official

D - Set Menu Editorial by en_translator


We assume that \(B\) is sorted in ascending order.

For each \(i\ (1 \leq i \leq N)\), consider how to find the sum of prices of set meals with the \(i\)-th main dish, i.e. \(\displaystyle \sum_{j=1}^{M} \min(A_i+B_j, P)\).

Let \(k_i\) be the maximum \(j\) such that \(B_j < P-A_i\) (or \(0\) if there is no such \(j\)). Then the sought value equals

\[ \begin{aligned} \displaystyle \sum_{j=1}^{M} \min(A_i+B_j, P) &= \sum_{j=1}^{k_i} (A_i+B_j) + \sum_{j=k_i+1}^{M} P \\ &=k_iA_i + \sum_{j=1}^{k_i} B_j + (M-k_i)P. \end{aligned} \]

The value of \(k_i\) can be found with a binary search, and \(\displaystyle \sum_{j=1}^{k_i} B_j\) as a cumulative sum.

Therefore, the problem has been solved in a total of \(O(N\log N)\) time.

Sample code (C++):

#include<bits/stdc++.h>

using namespace std;

using ll = long long;

int main() {
    int n, m, p;
    cin >> n >> m >> p;
    vector<int> a(n), b(m);
    for (int &i: a) cin >> i;
    for (int &i: b) cin >> i;
    sort(b.begin(), b.end());
    vector<ll> b_sum(m + 1);
    for (int i = 0; i < m; i++) {
        b_sum[i + 1] = b_sum[i] + b[i];
    }
    ll ans = 0;
    for (int i: a) {
        int lb = lower_bound(b.begin(), b.end(), p - i) - b.begin();
        ans += (ll) i * lb;
        ans += b_sum[lb];
        ans += (ll) p * (m - lb);
    }
    cout << ans << endl;
}

posted:
last update: