はまやんはまやんはまやん

hamayanhamayan's blog

StonesOnATree [SRM730 Div1 Easy]

https://community.topcoder.com/stat?c=problem_statement&pm=14811

頂点0を根とする木が与えられる。
木の各頂点は最大2個の子を含む。
頂点iには重みW[i]がある。
以下のルールに従って石を置く。
頂点0に石を置くための最大の瞬間重み総和は?

  • 子供(孫は関係ない)に石が置かれてないと、頂点に石は置けない
  • 自由に石を除いて良い
  • 葉は自由に石をおける
  • 石を置くとその頂点の重みが有効になる

解法

dfsで解いていこう。
木DPのようにしてやっていけばいい。
「dfs(cu) := 頂点cuに石を置くための最大瞬間重み総和」
これを計算していくが、子供の数が0~2なので、場合分けして処理していこう。
 
子供0なら葉なので、自分に石を置くだけ。
return W[cu]
 
子供1なら、子供に石を置いて(dfs(cu))、子供の子供の石を全て取り除いて自分に石を置く(W[cu]+W[child])。
瞬間最大なので、この2つの大きい方を返す
 
子供2なら、左側と右側のどちらを最初に石がある状態にするかで2通りあるので、それぞれやってみて最大瞬間重みの小さい方を採用する。
左側に最初に置く場合は、
左側の子供に石を置いて(dfs(left))
左側の子供の子供の石をすべて取り除いて、右側の子供に石を置く(dfs(right) + W[left])
右側の子供の子供の石をすべて取り除いて、自分に石を置く(W[cu] + W[left] + W[right])
瞬間最大なので、この3つの大きいものが最大瞬間重み。
右側もこれと同様にして、小さい方を返す。
 
Nが小さいので適用にやっても大丈夫そうだが、メモ化再帰しないとTLEで落ちる。

struct StonesOnATree {
    vector<int> E[1010];
    int N;
    vector<int> W;
    int memo[1010];

    int dfs(int cu) {
        if (0 < memo[cu]) return memo[cu];

        int res = inf;
        if (E[cu].size() == 0) res = W[cu];
        else if (E[cu].size() == 1) {
            res = max(W[cu] + W[E[cu][0]], dfs(E[cu][0]));
        } else {
            int lft = E[cu][0], rht = E[cu][1];

            // 左を作る -> 左の子供だけを残して右を作る -> 自分を作る
            chmin(res, max(W[cu] + W[lft] + W[rht], max(dfs(lft), W[lft] + dfs(rht))));

            // 逆
            chmin(res, max(W[cu] + W[lft] + W[rht], max(dfs(rht), W[rht] + dfs(lft))));
        }

        return memo[cu] = res;
    }

    int minStones(vector <int> p, vector <int> w) {
        N = w.size();
        W = w;
        rep(i, 1, N) E[p[i - 1]].push_back(i);
        return dfs(0);
    }
};