Mister雑記

競プロやります。

TopCoder SRM763 Medium 「RootItRight」

community.topcoder.com

問題

 N頂点の木が与えられる。頂点 vにはコスト V_vが割り当てられている。

この木の頂点を1つ選び rとし、木を rを根とする根付き木と見なす。 このとき、各頂点に対してコスト C_vを以下のように定める。

  • 頂点 vに対し、 rから vへのパスを r = v_0, v_1, \cdots, v_k = vとする。
  • このとき C_v = \sum_{i = 0}^{k} V_{v_i} \cdot i

そして根付き木のコストを、全頂点のコストの総和とする。

 rを適切に選んだときの、根付き木のコストの最小値を求めよ。

制約

  •  1 \leq N \leq 2 \times 10^5
  •  0 \leq V_i \lt 10^3

考察

根付き木に関する問題なので、見た感じ全方位木DPだろうな〜と検討がつく(実際合っている)。

まず根付き木のコストについてだが、これは以下のように変形することができる。

 \displaystyle \sum_{v = 0}^{N - 1} V_v \cdot d_v \cdot sz_v

ここで d_vは頂点 vの深さ、 sz_vは頂点 vを根とする部分木のサイズである。  vを通るパスは vの子孫の数だけあって、かつ vは常に d_v番目に現れることからこれは成り立つ。

これを元に根を変えたときのコストの変化量を考えると、以下の図のようになる。

よって頂点0を根としたときの

  • 頂点 vを根とする部分木のサイズ sz_v
  • 部分木全体の V_u \cdot sz_uの和 vsz_v

を事前計算で求めておけば良さそう。

ただし上の図にあるように部分木の根の扱いが厄介で、本番はこの辺を上手く詰められずに時間切れになった。

実装例

using lint = long long;

struct Edge { int src, dst; };
using Graph = vector<vector<Edge<Cost>>>;


class RootItRight {
    int VNUM;        // 頂点数
    vector<lint> V;  // 重み
    Graph tree;

    vector<lint> sz, vsz;
    lint ans;

    // sz, vszを埋めつつ初期解を求める
    // dは頂点vの深さ
    lint dfs(int v, int r, int d) {
        lint ret = 0;
        sz[v] = 1;
        vsz[v] = 0;
        for (auto e : tree[v]) {
            if (e.dst == r) continue;
            ret += dfs(e.dst, e.src, d + 1);
            sz[v] += sz[e.dst];
            vsz[v] += vsz[e.dst];
        }
        vsz[v] += sz[v] * V[v];
        ret += sz[v] * V[v] * d;
        return ret;
    }

    // vを根としたときの解がcost
    void edfs(int v, int r, lint cost, lint add) {
        ans = min(ans, cost);

        for (auto e : tree[v]) {
            if (e.dst == r) continue;

            lint nadd = add;
            nadd += vsz[v] - vsz[e.dst] - V[v] * sz[v];   // 根以外の頂点の分
            nadd += V[v] * (VNUM - sz[e.dst]);            // 根の分

            lint ncost = cost + nadd;
            ncost -= vsz[e.dst];

            edfs(e.dst, e.src, ncost, nadd);
        }
    }

public:
    lint findMinimumTotalCost(int N, vector<int> edge, vector<int> val, int D, int seed, int MX) {
        // グラフ、重みの初期化は省略
        sz.resize(N);
        vsz.resize(N);
        
        ans = dfs(0, -1, 0);

        edfs(0, -1, ans, 0);
        return ans;
    }
};

反省

  • DPは落ち着いて遷移を詰めないとダメとあれほど...。
  • グラフの入力形式に戸惑った時間が無駄だった。