Official

D - Set Menu Editorial by yuto1115

解説

以下、\(B\) が昇順にソートされていることを仮定します。

\(i\ (1 \leq i \leq N)\) について、\(i\) 種類目の主菜を用いた定食の価格の総和、すなわち \(\displaystyle \sum_{j=1}^{M} \min(A_i+B_j, P)\) を高速に求めることを考えます。

\(B_j < P-A_i\) を満たす最大の \(j\)(存在しない場合は \(0\))を \(k_i\) とおきます。このとき、求める値は、

\[ \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} \]

と表せます。\(k_i\) の値は二分探索を用いて、\(\displaystyle \sum_{j=1}^{k_i} B_j\) の値は累積和を用いて高速に計算可能です。

よって、本問題を \(O(N\log N)\) で解くことができました。

実装例 (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: