Codeforces Round #383 (Div. 2) D. Arpa's weak amphitheater and Mehrdad's valuable Hoses

問題

Problem - D - Codeforces

考察

問題を要約すると、「あるグループから"全員"・"1人"・"0人"のいずれかを選んでいったとき、重さがwを超えないbの合計値の最大を求めよ」となります。

まず、グループを判定するのはUnionFindで行えます(rootが同じなら同じグループ)。
各グループの合計重さ・美しさや構成員が判明したら、それを使って「グループごとに」DPを回せば答えが求まります。

DP票は2つ用意して交互に埋めてゆきます。
グループごとに回して交互に使うことで、同じグループの人が最大1回しか使われないようにすることができます。
また、合計は特殊な値として更新処理に加えてあげます。

コード

struct UnionFind {
    vector<int> data;

    UnionFind(int size) : data(size, -1) {}

    bool unionSet(int x, int y) {
        x = root(x);
        y = root(y);
        if (x != y) {
            if (data[y] < data[x]) swap(x, y);
            data[x] += data[y];
            data[y] = x;
        }
        return x != y;
    }

    bool findSet(int x, int y) {
        return root(x) == root(y);
    }

    int root(int x) {
        return data[x] < 0 ? x : data[x] = root(data[x]);
    }

    int size(int x) {
        return -data[root(x)];
    }
};

int main() {
    int n, m, mxw;
    cin >> n >> m >> mxw;

    vi w(n), b(n);
    rep(i, n)
        cin >> w[i];
    rep(i, n)
        cin >> b[i];

    UnionFind uf(n);
    rep(i, m) {
        int t1, t2;
        cin >> t1 >> t2;
        t1--;
        t2--;
        uf.unionSet(t1, t2);
    }

    vi dpw(n, 0), dpb(n, 0);
    map<int, vi> _group;
    rep(i, n) {
        int r = uf.root(i);
        dpw[r] += w[i];
        dpb[r] += b[i];
        _group[r].push_back(i);
    }
    vvi group;
    vector<pii> grp;
    tr(it, _group) {
        group.push_back(it->second);
        grp.push_back(mp(dpw[it->first], dpb[it->first]));
    }

    vvi memow(2, vi(mxw + 10, 0));
    vi* now = &memow[0];
    vi* next = &memow[1];
    rep(i, group.size()) {
        fill(all(*next), -1);
        rep(j, mxw + 5) {
            chmax((*next)[j], (*now)[j]);

            int no = (*now)[j];
            if(j + grp[i].first <= mxw) {
                chmax((*next)[j + grp[i].first], no + grp[i].second);
            }
            tr(it, group[i]) {
                if(j + w[*it] <= mxw) {
                    chmax((*next)[j + w[*it]], no + b[*it]);
                }
            }
        }
        swap(now, next);
    }

    int ans = 0;
    rep(i, mxw+1)
        chmax(ans, (*now)[i]);
    cout << ans << endl;
    
    return 0;
}

感想

グループ化の前処理をすることで(ほぼ)普通のDPとして扱えるようになるのは面白いなと思いました。あとタイトルが長いです。