Mister雑記

競プロ記事がメインになりますが、内容は気分次第です。

Codeforces Round #514 - E 「Split the Tree」 [2400]

codeforces.com

概要

頂点 1が根であるような n頂点根付き木が与えられる。各頂点には重みが割り振られていて、頂点 iの重みは w_iである。

この木を、以下の条件を満たすいくつかのパスによって分割する。

  • 1つの頂点はちょうど1つのパスに含まれる。
  • 各パスについて、頂点の重みの総和は S以下である。
  • 各パスに含まれる頂点数は L以下である。
  • パス上の頂点を v_1, v_2, \cdots, v_kとしたとき、 v_i v_{i-1}の親である。

このような分割が可能か判定し、可能ならばパスの数の最小値を求めよ。

制約

  •  1 \leq n \leq 10^5
  •  1 \leq L \leq 10^5
  •  1 \leq S \leq 10^{18}
  •  1 \leq w_i \leq 10^9

考察

まず「分割が不可能である」ことは「 \exists i \hspace{1ex} w_i \gt S」と同値になる。 これは全てのパスが1頂点になるような分割を考えれば分かる。


次に割り振り方だが、これは葉から貪欲に決めていくのが良さそうだなと直感的に思う。

各子について、子を根とする部分木をパスに分割したとする。こうして得られた分割において、その子を含むパスを「子のパス」と呼ぶことにする。このとき、

  • その頂点をいずれか1つの「子のパス」に含める。
  • 全ての子とその頂点を切り離す。

のうち最善のものを選んでいけばよい。

まず全ての「子のパス」が親を含む余裕がなければ、子は全て切り離さざるを得ない。一方で、もし余裕のある「子のパス」が存在すれば、その頂点はいずれかの「子のパス」へ含めるのが最適となる。

で、問題となるのは「どの子を選ぶか」。重みと頂点数という2つの指標があるため、これらの情報だけでは簡単には評価ができない。

解法

指標が2つあるのが悪いので、これらを1つにまとめてしまえばいい。

そこで各「子のパス」について、「あといくつの頂点を含めるか?」という値を考える。これが最も大きいものを貪欲に選べばいい。

「子のパス」の始点は決まっているので、始点から親へとパスを辿っていくシミュレーションによって、「パスに含むことができる頂点数」が計算できる。 これは愚直にやると計算量 O(n^2)だが、ダブリングをすることで O(n \log n)で行うことができる。


参考: betrue12.hateblo.jp

実装例

ダブリングはあまり実装したことがないので慣れていない。特に根をどう扱うかでとても悩んだ。

#include <iostream>
#include <vector>

using namespace std;
using ll = long long;

const ll INF = 1LL << 60;

int N, L;
ll S;
vector<int> path[100010];

ll upcost[100010][20];  // iを末端とする、長さ2^kのパスの重み
int upto[100010][20];   // iからの距離が2^kの祖先

int ans = 0;

// vを含む集合が、あといくつの要素を含めるか
int dfs(int v) {
    int up = 0;  // 子の値の最大値
    for (int sv : path[v]) up = max(up, dfs(sv));

    if (up > 0) {
        // 頂点vはこの集合に加える
        ans += path[v].size() - 1;
        return up - 1;

    } else {
        // 頂点vは新たな集合にする
        ans += path[v].size();

        int ret = 0;  // vを末端とする集合の最大可能サイズ
        ll cost = 0;
        while (v >= 0) {
            int k;
            for (k = 19; k >= 0; --k) {
                if (upcost[v][k] + cost <= S) break;
            }
            if (k < 0) break;  // 1つでも増やすとコストオーバー

            ret += (1 << k);
            cost += upcost[v][k];
            v = upto[v][k];
        }

        // サイズはL以下であることに注意
        return min(L, ret) - 1;
    }
}

int main() {
    cin >> N >> L >> S;

    for (int i = 0; i < N; ++i) {
        cin >> upcost[i][0];

        // Sを超える重みの頂点が1つでもあればアウト
        if (upcost[i][0] > S) {
            cout << -1 << endl;
            return 0;
        }
    }

    upto[0][0] = -1;
    for (int i = 1; i < N; ++i) {
        cin >> upto[i][0];
        path[--upto[i][0]].push_back(i);
    }

    // ダブリング
    for (int k = 1; k < 20; ++k) {
        for (int i = 0; i < N; ++i) {
            if (upto[i][k - 1] >= 0) {
                upto[i][k] = upto[upto[i][k - 1]][k - 1];
                upcost[i][k] = upcost[i][k - 1] + upcost[upto[i][k - 1]][k - 1];
            } else {
                // 0を超えて過剰に遡れないようにする
                upto[i][k] = -1;
                upcost[i][k] = INF;
            }
        }
    }

    dfs(0);
    // 頂点1の集合を加える
    cout << ans + 1 << endl;
    return 0;
}

感想

  • 2つの評価値を1つにまとめるという発想がなかった。
  • ダブリングの実装にも慣れていきたい。
  • ちなみにTutorialのDP解法はまるで理解できなかった。