はじめに

javascriptで再帰を使わずに全方位木DPするコードが必要になったので、せっかくだからTypeScriptで書いてみた。
参考記事は以下2つ。

コード

コード自体は @Kiri8128 氏のPython版のを参考にして、重み付き木用にアレンジして移植した代物です。

ReRooting.ts

class ReRooting<T> {
	N: number;
	edge: {to: number, cost: number}[][];
	merge: (left: T, right: T) => T;
	adjust: (acc: T, cost: number) => T;
	ident: T;
	constructor(N: number, merge: (left: T, right: T) => T, adjust: (acc: T, cost: number) => T, ident: T) {
		this.edge = Array.from({ length: N }, (v, k) => []);
		this.merge = merge;
		this.adjust = adjust;
		this.ident = ident;
		this.N = N;
	}
	add_edge(a: number, b: number, c: number) {
		this.edge[a].push({ to: b, cost: c });
		this.edge[b].push({ to: a, cost: c });
	}
	solve(): T[] {
		// Topological sort
		const topo = []; // ソート後のインデックスリスト
		// デフォルト値は根用のダミー。なおidx<0の場合はedgeを見ることはない
		const parent: { idx: number, edge: { to: number, cost: number } }[] = Array.from({ length: this.N }, (v, k) => ({ idx: -1, edge: { to: -1, cost: 0 } }));
		const q = [0];
		while (q.length > 0) {
			// undefinedになる場合があるとか文句言われたので形だけ回避。q.length>0だからそんなことはないのだけど
			let i = q.shift() ?? -1;
			topo.push(i);
			for (let e of this.edge[i]) {
				if (e.to == parent[i].idx) continue;
				parent[e.to] = { idx: i, edge: e };
				this.edge[e.to] = this.edge[e.to].filter(v => v.to != i);
				q.push(e.to);
			}
		}

		// Bottom-Up DP
		const dp = new Array(this.N);
		const bu = Array.from({ length: this.N }, (v, k) => this.ident);
		for (let i = this.N - 1; i > 0; i--) {
			const t = topo[i];
			let p = parent[t].idx;
			if (p >= 0) {
				dp[t] = this.adjust(bu[t], parent[t].edge.cost);
				bu[p] = this.merge(bu[p], dp[t]);
			}
		}
		dp[topo[0]] = bu[topo[0]];

		// Top-Down DP
		const td = Array.from({ length: this.N }, (v, k) => this.ident);
		for (let t of topo) {
			const es = this.edge[t];
			// 左からDP
			let ac = td[t];
			for (let e of es) {
				td[e.to] = ac;
				ac = this.merge(ac, dp[e.to]);
			}
			// 右からDP
			ac = this.ident;
			for (let j = es.length - 1; j >= 0; j--) {
				const { to, cost } = es[j];
				td[to] = this.adjust(this.merge(td[to], ac), cost);
				ac = this.merge(ac, dp[to]);
				dp[to] = this.merge(td[to], bu[to]);
			}
		}
		return dp;
	}
}

使い方は頂点数と適切なmerge, adjust, identを与えてlet solver = new ReRooting(N, merge, adjust, ident)でインスタンスを作って、solver.add_edge(from, to, weight)で重み付きの辺を登録した後solver.sovle()で計算結果を取得。
計算結果は各頂点を根とした場合のT型の値の配列なので、それも全部mergeしてsolver.solve().reduce(merge, ident)としてしまうことが多いかも。

使用例

最大経路長を求める場合はこんな感じ。

const merge = (a: number, b: number) => Math.max(a, b)
const adjust = (a: number, c: number) => a + c
const solver = new ReRooting(7, merge, adjust, 0)
solver.add_edge(0, 1, 3)
solver.add_edge(0, 2, 5)
solver.add_edge(1, 3, 1)
solver.add_edge(1, 4, 8)
solver.add_edge(2, 5, 2)
solver.add_edge(3, 6, 5)
console.log(solver.solve().reduce(merge, 0)) // 18

平均経路長を求める場合は、型Tを経路長の和と経路数の組にして、merge,adjustを次のようにすると良い。

const merge = (a: number[], b: number[]) => [a[0] + b[0], a[1] + b[1]]
const adjust = (a: number[], c: number) => [(a[1] + 1) * c + a[0], 1 + a[1]]
const solver = new ReRooting(7, merge, adjust, [0, 0])
solver.add_edge(0, 1, 3)
solver.add_edge(0, 2, 5)
solver.add_edge(1, 3, 1)
solver.add_edge(1, 4, 8)
solver.add_edge(2, 5, 2)
solver.add_edge(3, 6, 5)
const ret = solver.solve().reduce(merge, [0, 0])
console.log(ret[0] / ret[1]) // 8.857...