木が再帰的に定義されているときの重心分解

(* 番号づけられた頂点を持つ木 *)
(* T(i, v, ts) は根の番号が i、格納している値が v で、その下にある部分木のリストが ts である木を表す *)
type 'a indexed_tree = T of int * 'a * 'a indexed_tree list

let root_index (T(i, _, _)) = i
let root_value (T(_, v, _)) = v
let subtrees (T(_, _, ts)) = ts

(* 隣接リストから indexed_tree を構築する関数 *)
(* al はグラフの隣接リストで、al.(i) は i 番目の頂点と隣接する頂点の番号のリスト *)
let dfs_tree al =
  let rec dfs p i =
    T(i, (), List.filter_map (fun j -> if j = p then None else Some (dfs i j)) al.(i)) in
  dfs (-1) 0

(* 各頂点にそれを根とする部分木のサイズを格納した木を返す関数 *)
let rec size_tree t =
  let ts = List.map size_tree (subtrees t) in
  let count = List.fold_left (fun c t -> c + root_value t) 1 ts in
  T(root_index t, count, ts)

(* 重心分解を行う関数 *)
let centroid_decomp t =
  (*  t: いま注目している木 *)
  (* tt: t の上側にあった木のリスト *)
  (*  n: t と tt を合わせた木のサイズ *)
  (* Note: tt のサイズは常に 1 以下になります。そのため、list の代わりに option を使ってもいいですが、 *)
  (*       list の方がきれいに書けます。この小技は時々見ます *)
  let rec decomp t tt n =
    (* N.B. bigger のサイズは常に 1 以下 *)
    let (smaller, bigger) = List.partition (fun t -> root_value t * 2 <= n) (subtrees t) in
    match bigger with
    | t' :: _ ->
      let ts' = tt @ smaller in
      (* t' とそれ以外に分ける *)
      decomp t' [T(root_index t, n - root_value t', ts')] n in
    | [] ->
      T(root_index t, (), List.map (fun t -> decomp t [] (root_value t)) (tt @ subtrees t))
  decomp t [] (root_value t)

(* 木の各頂点の親を格納した配列を返す関数 *)
let parent_list_of_tree n t =
  let pl = Array.make n 0 in
  let rec dfs p (T(i, _, ts)) =
    pl.(i) <- p;
    List.iter (fun t -> dfs i t) ts in
  dfs (-1) t;
  pl

let () =
  let n = Scanf.scanf " %d" Fun.id in
  let al = Array.make n [] in
  let () =
    for i = 1 to n - 1 do
      Scanf.scanf " %d %d" (fun a b -> 
        (* 番号を 0 始まりにする *)
        let a = a - 1 in
        let b = b - 1 in
        al.(a) <- b :: al.(a);
        al.(b) <- a :: al.(b))
    done in
  let pl =
    al
    |> dfs_tree (* 隣接リストを indexed_tree に変換 *)
    |> size_tree (* すべての部分木のサイズを計算 *)
    |> centroid_decomp (* 重心分解 *)
    |> parent_list_of_tree n (* indexed_tree から各頂点の親の配列に変換 *) in
  for i = 0 to n - 1 do
    let p = pl.(i) in
    (* 番号を 1 始まりに戻す *)    
    let p = if p >= 0 then p + 1 else p in
    print_int p;
    if i < n - 1 then
      print_char ' '
    else
      print_newline()
  done

(* 各頂点を根とする部分木のサイズを格納した配列を返す関数 *)
(* al はグラフの隣接リストで、al.(i) は i 番目の頂点と隣接する頂点の番号のリスト *)
let subtree_size al =
  let ss = Array.make (Array.length al) 1 in
  let rec loop p i =
    List.iter (fun j ->
      if j <> p then begin
        loop i j;
        ss.(i) <- ss.(i) + ss.(j)
      end
    ) al.(i) in
  let () = loop (-1) 0 in
  ss

(* 重心分解を行う関数 *)
(* 戻り値 pl は重心分解で得られる木において i の親が pl.(i) であることを表す *)
(* al はグラフの隣接リスト、ss は subtree_size から得られる部分木のサイズを格納した配列 *)
let centroid_decomp al ss =
  let pl = Array.make (Array.length al) (-2) in
  (*  p: 直前に訪れた頂点 *)
  (*  i: いま訪れている頂点 *)
  (* pc: 1つ前に見つけた重心 *)
  (*  n: 全体の木のサイズ *)
  let rec decomp p i pc n =
    (* j <> p: 戻らない *)
    (* pl.(j) = -2: すでに見つけた重心は無視する *)
    let j = List.find_opt (fun j -> j <> p && pl.(j) = -2 && ss.(j) * 2 > n) al.(i) in
    match j with
    | Some j ->
      ss.(i) <- n - ss.(j);
      decomp i j pc n
    | None ->
      pl.(i) <- pc;
      List.iter (fun j -> if pl.(j) = -2 then decomp i j i ss.(j)) al.(i) in
  let () = decomp (-1) 0 (-1) (Array.length al) in
  pl

let () =
  let n = Scanf.scanf " %d" Fun.id in
  let al = Array.make n [] in
  let () =
    for i = 1 to n - 1 do
      Scanf.scanf " %d %d" (fun a b ->
        (* 番号を 0 始まりにする *)
        let a = a - 1 in
        let b = b - 1 in
        al.(a) <- b :: al.(a);
        al.(b) <- a :: al.(b))
    done in
  let pl = subtree_size al |> centroid_decomp al in
  for i = 0 to n - 1 do
    let p = pl.(i) in
    (* 番号を 1 始まりに戻す *)
    let p = if p >= 0 then p + 1 else p in
    print_int p;
    if i < n - 1 then
      print_char ' '
    else
      print_newline()
  done