LCA実装! 勝ち組へ

euler_tour/main.ts

@@ -2,19 +2,30 @@ export class EulerTour {
   #data: number[] = [];
   #l: number[] = [];
   #r: number[] = [];
-  constructor(edges: ([number, number]|[number, number, number])[], root = 0) {
+  #depth: number[] = [];
+  #cost: number[] = [];
+  constructor(
+    edges: ([number, number] | [number, number, number])[],
+    root = 0,
+  ) {
     const n = edges.length + 1;
-    const adj: number[][] = [];
+    const a: number[][] = [];
+    const c: number[][] = [];
     const prog: (0 | 1)[] = [];
     for (let i = 0; i < n; i += 1) {
-      adj.push([]);
+      a.push([]);
+      c.push([]);
       prog.push(0);
       this.#l.push(0);
       this.#r.push(0);
+      this.#depth.push(0);
+      this.#cost.push(0);
     }
     for (const edge of edges) {
-      adj[edge[0]].push(edge[1]);
-      adj[edge[1]].push(edge[0]);
+      a[edge[0]].push(edge[1]);
+      a[edge[1]].push(edge[0]);
+      c[edge[0]].push(edge[2] ?? 1);
+      c[edge[1]].push(edge[2] ?? 1);
     }
     const stack = [root];
     while (stack.length > 0) {
@@ -26,8 +37,12 @@ export class EulerTour {
       }
       this.#l[v] = this.#data.length - 1;
       prog[v] = 1;
-      for (const u of adj[v]) {
+      let i = -1;
+      for (const u of a[v]) {
+        i += 1;
         if (prog[u]) continue;
+        this.#depth[u] = this.#depth[v] + 1;
+        this.#cost[u] = this.#cost[v] + c[v][i];
         stack.push(v);
         stack.push(u);
       }
@@ -54,6 +69,24 @@ export class EulerTour {
     return this.#l[u] <= this.#l[v] && this.#r[v] <= this.#r[u];
   }
 
+  /**
+   * 頂点`v`の深さを返す
+   * @param v 対象の頂点
+   * @returns 根から頂点`v`へのパスに含まれる辺の数
+   */
+  depth(v: number): number {
+    return this.#depth[v];
+  }
+
+  /**
+   * 根から頂点`v`への距離を返す
+   * @param v 対象の頂点
+   * @returns 根から頂点`v`への距離
+   */
+  cost(v: number): number {
+    return this.#cost[v];
+  }
+
   *[Symbol.iterator]() {
     for (const v of this.#data) {
       yield v;

lca/main.ts

@@ -0,0 +1,67 @@
+import { EulerTour } from "../euler_tour/main.ts";
+
+export class LCA extends EulerTour {
+  #dup: number[][] = [];
+  constructor(
+    edges: ([number, number] | [number, number, number])[],
+    root = 0,
+  ) {
+    const len = edges.length + 1;
+    super(edges, root);
+    let p: number[] = [];
+    for (let i = 0; i < len; i += 1) {
+      p.push(this.parent(i));
+    }
+    this.#dup.push(p);
+    if (len === 1) return;
+    while (1) {
+      const np: number[] = [];
+      let count = 0;
+      for (let i = 0; i < len; i += 1) {
+        const v = p[p[i]] ?? -1;
+        if (v !== -1) count += 1;
+        np.push(v);
+      }
+      if (count === 0) break;
+      this.#dup.push(np);
+      p = np;
+    }
+    this.#dup.reverse();
+  }
+
+  /**
+   * 頂点`u`と頂点`v`の最小共通祖先を求める
+   * @param u 1つめの頂点の番号
+   * @param v 2つめの頂点の番号
+   * @returns 最小共通祖先の番号
+   */
+  lca (u: number, v: number): number {
+    if (this.inSubtree(u, v)) return u;
+    let y = u;
+    for (const p of this.#dup) {
+      if (p[y] === -1) continue;
+      if (!this.inSubtree(p[y], v)) y = p[y];
+    }
+    return this.parent(y);
+  }
+
+  /**
+   * 根から頂点`v`への距離を返す
+   * @param v 対象の頂点
+   * @returns 根から頂点`v`への距離
+   */
+  cost (v: number): number;
+  /**
+   * 頂点`u`から頂点`v`への距離を返す
+   * @param u 1つめの頂点
+   * @param v 2つめの頂点
+   * @returns 頂点`u`,`v`間の距離
+   */
+  cost (u: number, v: number): number;
+  cost (u: number, v?: number): number {
+    if (!v) { // vは0 or undefined.どちらも根からの距離
+      return super.cost(u);
+    }
+    return super.cost(u) + super.cost(v) - 2 * super.cost(this.lca(u, v));
+  }
+}