これは楽しい！！！
順番をうまく決めれば、全部の情報を持たなくても、個数のみの情報に落とせる系。

問題概要

$N$ 要素の数列 $a_{1}, \dots, a_{N}$ と、 $N \times P$ の二次元数列 $s$ が与えられる。

$N$ 個の $a$ の中から $K$ 個を選び、
残りの $N-K$ 個の index の中から $P$ 個分、 $s$ の行ベクトルを選び、どの列も共有しないように列の index を選ぶ

こうして得られる合計 $K + P$ 個の整数の総和の最大値を求めよ。

制約

$2 \le N \le 10^{5}$
$1 \le K + P \le N$
$1 \le P \le 7$

考えたこと

要は $N \times (P+1)$ の長方形の各列について、どの 2 つの列も横方向に並ばないようにしつつ

最終列以外は 1 個の要素を選び
最終列は $K$ 個の要素を選び

総和を最大化する問題となる。 $P$ が小さいので bitDP ができる。難しいのは $K$ はとても大きくなりうるので、以下のような素朴な bitDP では間に合わない。

dp[ i ][ bit ][ num ] := i 行目まで見て、最終列以外は bit に対応する列の index をすでに選んでいて、最終列は num 個の要素をすでに選んでいる場合についての最大値

この bit DP は自然だけど、 $O(NP2^{P}K)$ かかってしまう。

最終列をソート

最終列以外についてどの行の要素を選ぶか決めると、最終列については残った行のうち、値が大きい順に $K$ 個選べば良いことがわかる。このことから、

最終列の値が大きい順にソートしておくと
i 個目の要素を見ているときに、bit のみの情報から、最終列の要素をとるべきかどうかが決まる

ということがわかる。よって、DP の情報量を落として

dp[ i ][ bit ] := i 行目まで見て、最終列以外は bit に対応する列の index をすでに選んでいる状態についての最大値

として解くことができる。これなら計算量は $O(NP2^{P})$ で済む。

#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 N, P, K;
vector<long long> a;
vector<vector<long long>> s;

long long solve() {
    vector<int> ids(N);
    iota(ids.begin(), ids.end(), 0);
    sort(ids.begin(), ids.end(), [&](int i, int j) {
            return a[i] > a[j];});
    vector<vector<long long>> dp(N+1, vector<long long>(1<<P, -INF));
    dp[0][0] = 0;
    for (int i = 0; i < N; ++i) {
        for (int bit = 0; bit < (1<<P); ++bit) {
            int con = __builtin_popcount(bit);
            if (i - con < K) chmax(dp[i+1][bit], dp[i][bit] + a[ids[i]]);
            else chmax(dp[i+1][bit], dp[i][bit]);
            for (int j = 0; j < P; ++j) {
                if (bit & (1<<j)) continue;
                int nbit = bit | (1<<j);
                chmax(dp[i+1][nbit], dp[i][bit] + s[ids[i]][j]);
            }
        }
    }
    return dp[N][(1<<P)-1];
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0); 

    cin >> N >> P >> K;
    a.resize(N);
    s.assign(N, vector<long long>(P, 0));
    for (int i = 0; i < N; ++i) cin >> a[i];
    for (int i = 0; i < N; ++i) for (int j = 0; j < P; ++j) cin >> s[i][j];
    cout << solve() << endl;
}

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

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

Codeforces CodeCraft-20 (Div. 2) E. Team Building (R2300)

問題概要

制約

考えたこと

最終列をソート