はじめに
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...