【ABC216】F - Max Sum Counting

atcoder.jp

解法

愚直に解こうとするとO(2^N)で絶対間に合わない。

Aはmax、Bは総和が条件として示されているため、A[i] ≦ 5000という制約が気になる。
この制約があることで、Bの和が5000より大きい数になるような集合は考えなくて良いことになる。
Bの要素を適当に選んで、5000までの和の作り方が何通りか、という問題は下のようなDPで解くことができる。
dp[i][j] ← i番目までの要素を見たときに和がjであるような選び方は何通りか
これは、部分和問題(またはナップザック問題)と同じである。

さらにAはmaxが条件になっていることから、なんとなく主客転倒の考え方が使えそうな気がしてくる。
そこでAを昇順に並び替えて、集合にいれることができるのは既に見た要素のみに制限し、前から順に要素を見ていくことを考えてみる。
ここで、BもAに合わせて並び替える必要があることに注意する。
昇順に見ていくため、今見ている要素を選ぶときは、必ず集合の最大がその要素になることを利用することができる。

i番目の要素が条件を満たすには、次の2つの条件を同時に満たすことが必要になる。
(Aを昇順に並び替えた時のi番目の要素をA[i], B[i]と表すことにする)

• B[i]を集合に含む
• Bの集合の和がA[i]を超えない

この条件を満たす集合の選び方がいくつかを、上のDPを解きながら数えていけば良い。

実装

void solve(long long N, std::vector<long long> A, std::vector<long long> B) {

    // 順序を保ってAを基準に昇順ソート
    vector<pair<ll, ll>> vec;
    REP (i, N) vec.push_back({A[i], i});
    ALL(sort, vec);

    mint ans = 0;
    // dpテーブル
    vector<vector<mint>> cnt(N+1, vector<mint>(5001, 0));
    cnt[0][0] = 1;

    // 上の解説のA[i]はvec[i].first、B[i]はvec[i].secondになっている
    REP (i, N) {
        REP (j, 5001) {
            cnt[i + 1][j] += cnt[i][j]; // i番目の要素を選ばないとき
            // i番目の要素を選ぶとき
            if (j + B[vec[i].second] <= 5000) {
                cnt[i + 1][j + B[vec[i].second]] += cnt[i][j];
                if (vec[i].first >= j + B[vec[i].second]) {
                    // 答えに加算するのはi番目の要素を選び、A[i](ここではvec[i].first)が集合の総和を超えないとき
                    ans += cnt[i+1][j + B[vec[i].second]];
                }
            }
        }
    }

    c(ans); // 出力
}