Submission #230665

---

Source Code Expand

let pf = Printf.printf ;;
let epf = Printf.eprintf ;;
let sf = Scanf.scanf ;;

let (|>) x f = f x ;;
let (@@) f x = f x ;;

module Array = struct
  include ArrayLabels
  let foldi ~f ~init arr =
    let acc = ref init in
    for i = 0 to Array.length arr - 1 do
      acc := f i !acc arr.(i)
    done;
    !acc
  ;;

  let findi arr ~f =
    let len = Array.length arr in
    let rec iter i =
      if i = len then None
      else if f arr.(i) then Some i
      else iter (i + 1) in
    iter 0
  ;;

  let findi_exn arr ~f =
    let len = Array.length arr in
    let rec iter i =
      if i = len then None
      else if f arr.(i) then Some i
      else iter (i + 1) in
    iter 0
  ;;

  (* i ∈ [0..res) -> not (f i) /\ i ∈ [res, len) -> f i *)
  let lower_bound arr ~f =
    let l, h = 0, Array.length arr in
    (* res ∈ (l, h] *)
    let rec iter l h =
      if l = h - 1 then h
      else 
        let m = (l + h) lsr 1 in
        if f m then iter l m
        else iter m h in
    iter l h
  ;;

  (* i ∈ [0..res] -> f i /\ i ∈ (res, len) -> not (f i) *)
  let upper_bound arr ~f =
    let l, h = 0, Array.length arr in
    (* res ∈ [l, h) *)
    let rec iter l h =
      if l = h - 1 then l
      else 
        let m = (l + h) lsr 1 in
        if f m then iter m h
        else iter l m in
    iter l h
  ;;
end
  
module String = StringLabels ;;
module List  = struct
  include ListLabels ;;
  let rec repeat n a = if n = 0 then [] else a :: repeat (n - 1) a ;;
  let rec drop n a =
    if n = 0 then a
    else match a with
    | [] -> failwith "cannot take"
    | x :: xs -> drop (n - 1) xs ;;

  let init ~f n =
    let res = ref [] in
    for i = 0 to n - 1 do
      res := f i :: !res
    done;
    List.rev !res
  ;;
end ;;
module H = Hashtbl ;;

let c = ref ' ' ;;
let is_num c = '0' <= c && c <= '9' ;;
let is_space c = c = ' ' || c = '\n' || c = '\t' ;;
let to_num c = Char.code c - Char.code '0' ;;
let read_char () = input_char stdin ;;

let next_int () =
  while is_space !c do
    c := read_char () ;
  done;
  let ok = ref false in
  let acc = ref 0 in
  while is_num !c do
    ok := true;
    acc := !acc * 10 + to_num !c;
    c := read_char ();
  done;
  if !ok then !acc
  else raise (Failure "next_int")
;; 

module SI = Set.Make (struct
  type t = int 
  let compare = compare
end)

let solve n graph q query =
  let parent = Array.create_matrix 30 n 0 in
  let depth = Array.create n 0 in
  
  let rec dfs v p d =
    parent.(0).(v) <- p;
    depth.(v) <- d;
    List.iter graph.(v) ~f:(fun next ->
      if next <> p then dfs next v (d + 1)) in
    
  let init () =
    dfs 0 (-1) 0;
    for k = 0 to 20 do
      for v = 0 to n - 1 do
        if (parent.(k).(v) < 0) then parent.(k + 1).(v) <- -1
        else parent.(k + 1).(v) <- parent.(k).(parent.(k).(v))
      done
    done in

  let lca u v =
    let u, v = if depth.(u) > depth.(v) then ref v, ref u else ref u, ref v in
    let acc = ref (depth.(!v) - depth.(!u)) in
    for k = 0 to 20 do
      if ((depth.(!v) - depth.(!u)) lsr k) land 1 <> 0 then
        v := parent.(k).(!v)
    done;
    
    if !u = !v then !u, !acc
    else begin
      for k = 20 downto 0 do
        if parent.(k).(!u) <> parent.(k).(!v) then begin
          u := parent.(k).(!u);
          v := parent.(k).(!v);
          acc := !acc + (1 lsl k) * 2
        end
      done;
      parent.(0).(!u), (!acc + 2)
    end in

  init ();
  for i = 0 to q - 1 do
    let a, b = query.(i) in
    let (v, x) = lca a b in
    pf "%d\n" (x + 1)
  done
;;

let () =
  let n = next_int () in
  let graph = Array.create n [] in
  for i = 1 to n - 1 do
    let a = next_int () in
    let b = next_int () in
    let a, b = a - 1, b - 1 in
    graph.(a) <- b :: graph.(a);
    graph.(b) <- a :: graph.(b);
  done;
  let q = next_int () in
  let query = Array.init q ~f:(fun _ ->
    let a = next_int () in
    let b = next_int () in
    a - 1, b - 1) in
  solve n graph q query
;;

Submission Info

Judge Result