けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

Codeforces Round #681 (Div. 1) D. Sum (R2800)

 O(NK^{2}) から落とせる気がまったくしなかった...

問題概要

 N 個の広義単調増加数列  a_{i,1}, a_{i,2}, \dots, a_{i,t_{i}} が与えられる。

それぞれの数列から、先頭から  k_{i} 個ずつとってきた値の総和の最大値を求めよ。ただし  k_{1} + \dots, k_{N} = K でなければならないものとする。

制約

  •  1 \le N, K \le 3000
  •  1 \le \sum_{i=1}^{N}t_{i} \le 10^{6}

考えたこと

個数に関する convolution の形をした DP の問題。ナイーブには  O(NK^{2}) でできるが...

とりあえず、「各数列が広義単調増加」という怪しすぎる条件について考察する。これを使うと最終結果において

ただ 1 個の数列を除いて、他の数列は「まったくとらない」「すべてとる」のいずれか (中途半端なのは 1 個だけ)

になるということが言える。ある解が存在して最適だったとき、中途半端な個数の数列が 2 つ以上あったとき、それらのどちらかを伸ばしてどちらかを下げることで解を改善できることが言えるからだ。

O(NK2) から O(N2 K) へ

上記の知見を活かすと、次のような解法が生える。

  • 中途半端な数列を全探索する ( O(N) 倍)
  • それ以外の  N-1 個の数列については「まったくとらない」「すべてとる」のいずれかなので、「 f(k) = 全体で  k 個採用したときの最適値」として、 O(NK) f(k) が求められる

この解法の計算量は  O(N^{2}K) となる。ここまでで計算量は、先ほどの解法も合わせると  O(NK \min(N, K)) にはなった。でもまだ足りない。

このような  N-1 個の結果をそれぞれ求める方法としては

  • 左右からの累積結果をあらかじめ求めておく
  • 戻す DP

あたりが有力。しかし 1 個目については、最後にまとめるときに「 k 方向の畳み込み」が必要となって  O(NK^{2}) の計算量となってしまう。2 個目については、今回の DP 更新は不可逆なので戻せない。

分割統治法

 N-1 個についての結果を求める第三の方法として、分割統治法がある。考えてみれば、たとえば  N = 16 として、下図の 2 つの「赤い箇所」に対応する結果を計算するとき、「黄色い箇所」に関する結果は使いまわせそう。

f:id:drken1215:20201103162209p:plain

このような考え方を洗練させると、分割統治法に基づく  O(KN\log N) の解法が生える。

コード

#include <bits/stdc++.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

const long long INF = 1LL<<60;
int main() {
    cin.tie(0); 
    ios::sync_with_stdio(false);

    int N, K;
    cin >> N >> K;
    vector<vector<long long>> A(N), S(N);
    vector<long long> T(N);
    for (int i = 0; i < N; ++i) {
        cin >> T[i];
        A[i].resize(T[i]);
        S[i].assign(T[i]+1, 0);
        for (int j = 0; j < T[i]; ++j) cin >> A[i][j], S[i][j+1] = S[i][j] + A[i][j];
    }

    auto add = [&](int id, vector<long long> &dp) -> void {
        for (int k = K; k >= T[id]; --k) chmax(dp[k], dp[k-T[id]] + S[id].back());
    };
    long long res = -INF;
    auto rec = [&](auto self, int left, int right, const vector<long long> &dp) -> void {
        if (right - left <= 1) {
            for (int k = 0; k <= K; ++k) {
                if (K-k < S[left].size()) chmax(res, dp[k] + S[left][K-k]);
            }
            return;
        }
        int mid = (left + right) / 2;
        // left
        {
            auto dp2 = dp;
            for (int pos = mid; pos < right; ++pos) add(pos, dp2);
            self(self, left, mid, dp2);
        }
        // right
        {
            auto dp2 = dp;
            for (int pos = left; pos < mid; ++pos) add(pos, dp2);
            self(self, mid, right, dp2);
        }
    };

    vector<long long> dp(K+1, -INF);
    dp[0] = 0;
    rec(rec, 0, N, dp);
    cout << res << endl;
}