AtCoder精選良問「D - Ki」Diff 920

この記事はAtCoderの精選良問集マガジンの中の記事です。
私が良問だと思った問題を解説付きでまとめています。

https://note.com/kuri_tter/m/m622eeb3d9ca2

問題リンク

D - Ki

問題概要

はじめ、１からNまでの番号が振られたN頂点の根付き木が与えられる。
i番目の辺は頂点a_iと頂点b_iを結んでいる。
各頂点にはカウンターが設置されている。
下記のようなQ回の操作を行う。
頂点pを根とする部分木に含まれるすべての頂点のカウンターの値にxを足す。
すべての操作後の各頂点のカウンターの値を答える、という問題

考え方

まずは入力例１の図を見てみる。根付き木とは以下のような親子関係のあるグラフ構造。

行う操作は要は以下のようになる。
「指定された親頂点とその全ての子供のカウンターをX足す」
上から水を流すイメージをするといいかもしれない。

ちなみに私はタイトルが「Ki」とあるので、なんとなく「気」のようなものが下に流れるイメージも持った。作問者が「木」とかけたダジャレを放っていたのかは分からないが、私は気の利いたダジャレだと思った。ちなみに「気の利いた」という言葉もダジャレである。

問題に話を戻そう。
制約はNが100,000でQも100,000。
操作が与えられるたびに足し上げていっても答えは出せるが、制約的にその都度やっていたら明らかに間に合わない。

そのため、なんとか一発で答えを求める方法を考える。

まずは各クエリの都度に探索をするのではなく
{ 頂点1：足される量, 頂点2：足される量… }
というような辞書を用意する方針でいこう。

そのうえでグラフを１番から探索していくことを考えて、幅優先か深さ優先のどちらで攻めるかを考える。

結論を言うと深さ優先の方を選択するのが正しい。

幅優先で探索をしていくと上の図の場合、以下のような順番でたどっていく。

1 => 2 => 3 => 4

なんとなく３と４に枝わかれしてしまっているので、２に水を流した後に３に水を流すと計算がややこしくなりそうな気がする。順番に足し上げていくと、４に流れる量が正しいものよりも多くなってしまう。

この「ややこしくなりそう」という直感をひとまず信じて、深さ優先の順番をみてみる。深さ優先の場合は以下のような順番でたどる。

1 => 2 => 3 => 2 => 4

深さ優先探索は猪突猛進的に掘り進むイメージ。
言い換えると今回のケースで必ず「親=>子」という順番に進む。
先にすすめる子がいなくなったら「子=>親」に戻る、という感じだ。

こちらの順番なら解ける気がする。
各クエリのデータは以下のような辞書にまとめていた。

{ 頂点1：足される量, 頂点2：足される量… }

まずは１番の親から頂点を順番にたどっていって、流れる水の量をどんどん足していけばいい。親に流れた水は必ずその子供に流れていくからだ。

そうして「子=>親」の遷移をした際は、子に流した水を減らすような感じで流量を調整する。
そのような実装をする。

実装

import { createInterface } from "readline";
import * as fs from "fs";
import { log } from "console";
import { exit } from "process";


function main() {
    const [n, q] = nextNums(2)
    // graphを構築
    const graph: Record<number, number[]> = []
    for (let i = 0; i < n - 1; i++) {
        const [a, b] = nextNums(2)
        if (graph[a]) {
            graph[a].push(b)
        } else {
            graph[a] = [b]
        }
        if (graph[b]) {
            graph[b].push(a)
        } else {
            graph[b] = [a]
        }
    }
    // 各頂点に流れる量を保持
    const flowDict: Record<number, number> = {}
    // 0で初期化
    for (let i = 1; i < n + 1; i++) {
        flowDict[i] = 0
    }
    for (let i = 0; i < q; i++) {
        const [a, x] = nextNums(2)
        flowDict[a] += x
    }
    // 訪問済みの頂点を記録
    const vis: Set<number> = new Set()
    vis.add(1)
    // 答えの配列
    const answer: number[] = new Array(n + 1).fill(0)
    // 流れている量
    let flow = 0
    // 深さ優先探索
    const stack: number[] = [1]
    while (stack.length > 0) {
        const now = stack.pop()!
        if (now > 0) {
            // 水を足す
            flow += flowDict[now]
            // 答えに足す
            answer[now] += flow
            if (!graph[now]) continue
            // 子供を見ていく
            for (const child of graph[now]) {
                // 訪問済みはスキップ
                if (vis.has(child)) continue
                vis.add(child)
                // マイナスを先に入れる
                stack.push(-1 * child)
                stack.push(child)
            }
        } else {
            // マイナスの頂点が返ってきた
            // つまり次は子=>親の遷移になる
            // あげた水を減らす
            flow -= flowDict[-1 * now]
        }
    }
    // 答えの出力。１インデックスにしているためsliceする。
    log(answer.slice(1).join(' '))
}

let inputs = "";
let inputArray: string[];
let currentIndex = 0;

function next() {
    return inputArray[currentIndex++];
}
function nextNum() {
    return +next();
}
function nextBigInt() {
    return BigInt(next());
}
function nexts(length: number) {
    const arr = [];
    for (let i = 0; i < length; ++i) arr[i] = next();
    return arr;
}
function nextNums(length: number) {
    const arr = [];
    for (let i = 0; i < length; ++i) arr[i] = nextNum();
    return arr;
}
function nextBigInts(length: number) {
    const arr = [];
    for (let i = 0; i < length; ++i) arr[i] = nextBigInt();
    return arr;
}


// デバッグ環境がWindowsであれば条件分岐する
if (process.env.OS == "Windows_NT") {
    const stream = createInterface({
        input: process.stdin,
        output: process.stdout,
    });
    stream.on("line", (line) => {
        inputs += line;
        inputs += "\n";
    });
    stream.on("close", () => {
        inputArray = inputs.split(/\s/);
        main();
    });
} else {
    inputs = fs.readFileSync("/dev/stdin", "utf8");
    inputArray = inputs.split(/\s/);
    main();
}