TopCoder SRM 730 Div1 Easy (Div2 Hard): StonesOnATree

するめうめ

tsutaj.hatenablog.com

問題概要

原文 → TopCoder Statistics - Problem Statement

頂点に重み がついている、 頂点からなる木がある。子よりも親の方がインデックスが小さく、重みについても子と同じか小さいことが保証される。

この木に石を置いたり、取り除いたりすることを考える。ある頂点に石が置ける条件とは、その頂点の子すべてに石が置かれていることである。取り除く際は、どこを取り除いてもよい。

ある状態に対して、石が置かれている頂点に対応する重みの和をコストとするとき、頂点 に石を置くために必要な最大コストの最小値を求めよ。

解説

ドワコン 2018 予選の D 問題 と設定と似ていますが、「重みが単調増加である」という点が異なります。

かつ であることから、任意の頂点 の重みについて、その子の重みと等しいか小さくなります。この制約により、ある頂点 に対して、「その子 に石が置けるように操作し、それが終わってから別の子 に石が置けるように操作する」ことが最適になります。 に対して動かしている途中に に対して動かす必要がないということです。

よって、 頂点 に石を置くために必要な最大コストの最小値 とし、これを葉から根の方向に更新すれば良いです。

Div1 Easy では子の数は高々 なので、場合分けします。

子の数が である場合は、葉なので明らかに が入ります。

である場合は、自分の重みとその子の重みの和と、子に石を置くために必要なコストの です。

である場合について考えます。説明のため子である頂点を とします。このとき、子に石を置く順番は または のどちらかです。このうち、コストが小さくなる方を選択すれば良いです。

Div2 Hard では二分木とは限らないので、もう少し考えます。

頂点 の子を とします。とりあえず、 の順番に石を置き、それから に石を置くことにしましょう。

を置くまでのコストは、 DP で計算している値そのものなので、 です。

を置くまでのコストは、既に が置かれているため になります。

このように、石を置くのが 番目となる子 について、そのコストは となり、あとは適切に石を置く順番を決めることでコストを最小化しようということになります。

ここで、任意の について、 であることを利用します (これは を置くために必ず かかることを考慮すれば明らかです)。 ですから、 の値と の値の差分が大きい物から先に配置していけば、最終的なコストは小さくなります。よって、このルールに従って子をソートすれば解けます。

最後に、子の重みの総和と自分自身の重みの和を考慮し忘れないように気をつけましょう。

ソースコード

Div1

class StonesOnATree {
    public:
    int minStones(vector<int> p, vector<int> w) {
        int N = w.size();
        vector< vector<int> > children(N);
        for(int i=0; i<N-1; i++) {
            children[ p[i] ].push_back(i+1);
        }

        vector<int> dp(N);
        for(int i=N-1; i>=0; i--) {
            int sz = children[i].size();
            if(sz == 0) {
                dp[i] = w[i];
            }
            if(sz == 1) {
                int u = children[i][0];
                dp[i] = max(dp[u], w[i] + w[u]);
            }
            if(sz == 2) {
                int u = children[i][0], v = children[i][1];
                // u -> v
                int vl = max(dp[u], w[u] + dp[v]);
                // v -> u
                int vr = max(dp[v], w[v] + dp[u]);
                dp[i] = max(w[i] + w[u] + w[v], min(vl, vr));
            }
        }
        return dp[0];
    }
};

Div2

class StonesOnATreeDiv2 {
    public:
    int minStones(vector<int> p, vector<int> w) {
        int N = w.size();
        vector< vector<int> > children(N);
        for(int i=0; i<N-1; i++) {
            children[ p[i] ].push_back(i+1);
        }

        vector<int> dp(N);
        for(int i=N-1; i>=0; i--) {
            sort(children[i].begin(), children[i].end(), [&](int l, int r) {
                return dp[l] - w[l] > dp[r] - w[r];
            });

            int sum_weight = 0;
            for(size_t k=0; k<children[i].size(); k++) {
                int u = children[i][k];
                dp[i] = max(dp[i], sum_weight + dp[u]);
                sum_weight += w[u];
            }
            sum_weight += w[i];
            dp[i] = max(dp[i], sum_weight);
        }
        return dp[0];
    }
};