くれなゐの雑記

例を上げて 自分で手を動かして学習できる入門記事を多めに書いています

Educational Codeforces Round 54 E. Vasya and a Tree

問題

codeforces.com

問題概要

Tree(重み付きではない)が与えられ、木の各頂点に数字を書き込む。最初は0. 以下のM個のクエリも与えられる。

v, d, x: 頂点vと、その部分木のうち、vからの距離がd以内に含まれる頂点にxを足す。

昔似た方針で解いた記憶がある。リアルタイムimosって感じ。

DFSで頂点を見ていく。DFSでは、一個前の頂点で書き込んだ数字を引数として持っておく(sとする)。 見た頂点から、d,xなクエリが発行されていたら、sxを一通り足す。 同時に、木の最大の深さ分の配列を用意しておいて、今見ている頂点の深さに、dを足した先にxを足してメモしておく。 式で表現すると、今見ている頂点の深さをdepthとすると

memo[depth+d] += x

みたいな感じ。xを一通り足したら、memoでメモされている分を引く。(そのクエリの範囲外に到達するため)

また、今見ている頂点が見終わったら、memoに足した分は無効化されるので、最後に引いて終了しないといけない。(ソースコード参照)

vector<vector<int> > G;
vector<vector<pii> > dist_value;
vector<int> res;
vector<int> imos;

vector<bool> is_visited;
void DFS(int node, int depth, int sum){
    if (is_visited[node]) return;
    is_visited[node] = true;

    for (auto &dx : dist_value[node]) {
        int d, x; tie(d, x) = dx;
        int right = min(size_t(depth + d + 1), G.size());
        imos[right] -= x;
        sum += x;
    }

    sum += imos[depth];
    res[node] = sum;

    for (auto &dst : G[node]) {
        DFS(dst, depth + 1, sum);
    }

    for (auto &dx : dist_value[node]) {
        int d, x; tie(d, x) = dx;
        int right = min(size_t(depth + d + 1), G.size());
        imos[right] += x;
    }
}

void solve() {
    int N; cin >> N;
    G.resize(N);
    dist_value.resize(N);
    res.resize(N);
    imos.resize(N+1);
    is_visited.resize(N);

    REP(i, N-1) {
        int x, y; cin >> x >> y;
        --x, --y;
        G[x].push_back(y);
        G[y].push_back(x);
    }

    int M; cin >> M;
    REP(i, M) {
        int v, d, x; cin >> v >> d >> x;
        --v;
        dist_value[v].emplace_back(d,x);
    }

    DFS(0, 0, 0);

    REP(i, res.size()) {
        cout << res[i] << (res.size() == i + 1 ? '\n' : ' ');
    }
}