OCaml で競技プログラミングをやっている時の知見

module Modulo (M : sig include Int_intf.S val modulo : t end) = struct
  open M

  let rec inverse ?(b = modulo) ?(u = one) ?(v = zero) a =
    if b = zero then (u % modulo + modulo) % modulo
    else
      let t = a / b in
      let a, b = b, a - t * b in
      let u, v = v, u - t * v in
      inverse ~b ~u ~v a
  let ( / ) a b = a * inverse b

  let ( + ) a b = (a % modulo + b % modulo) % modulo
  let ( * ) a b = (a % modulo * b % modulo) % modulo

  let power a b =
    Array.init num_bits ~f:Fn.id
    |> Array.fold ~init:(one, a) ~f:(fun (p, q) i ->
      if b land (one lsl i) <> zero then p * q, q * q
      else
        p, q * q
    )
    |> Tuple2.get1
  let ( ** ) a b = power (a % modulo) (b % modulo)
end

module Modulo998244353 = Modulo (struct include Int let modulo = 998244353 end)

  let ( + ) a b = (a % modulo + b % modulo) % modulo
  let ( * ) a b = (a % modulo * b % modulo) % modulo

  let comb n r = 
    let rec loop z i =
      if r < i then z
      else loop ((n - r + i) / i * z) (i + one)
    in
    loop one one

let (--) start end_ k =
  for i = start to end_ do
    k i
  done
let (let+) seq f k = seq (fun x -> f x k)
let (let*) seq f k = seq (fun x -> k (f x))
let (let^) seq p k = seq (fun x -> if p x then k x)

let (.!())   = Graph.around
let (.!()<-) = Graph.push
let (.%()<-) = Graph.connect

module DirectedGraph = struct
  type t = {
    size : int;
    normal  : (int, int Iter.t) Hashtbl.t;
    reverse : (int, int Iter.t) Hashtbl.t;
  }
  let create ~size = {
    size;
    normal  = Hashtbl.create ~size (module Int);
    reverse = Hashtbl.create ~size (module Int);
  }

  let push v = function
    | None     -> Iter.singleton v 
    | Some acc -> Iter.snoc acc v
  let push g v u = Hashtbl.update g v ~f:(push u)
  let push g v u = push g.normal v u; push g.reverse u v

  let next g v = Hashtbl.find g.normal  v |> Option.value ~default:Iter.empty
  let pred g v = Hashtbl.find g.reverse v |> Option.value ~default:Iter.empty

  let strongly_connected_components g =
    let iterate visit f = 
      Iter.fold 
        (fun acc v -> if visit.(v) then acc else Iter.cons (f v) acc)
        Iter.empty
    in
    let step1 = 
      let visit = Array.init (g.size + 1) ~f:(const false) in
      let rec dfs acc v =
        visit.(v) <- true;
        next g v
        |> Iter.filter (fun u -> not visit.(u))
        |> Iter.fold dfs acc
        |> Iter.cons v
      in
      iterate visit @@ dfs Iter.empty
    in
    let step2 = 
      let visit = Array.init (g.size + 1) ~f:(const false) in
      let rec dfs acc v =
        visit.(v) <- true;
        pred g v
        |> Iter.filter (fun u -> not visit.(u))
        |> Iter.fold dfs (Iter.cons v acc)
      in
      iterate visit @@ dfs Iter.empty
    in
    Iter.(1 -- g.size) |> step1 |> Iter.flatten |> step2
end

let n = scanf "%d" Fn.id
let a = Array.init n ~f:(fun _ -> scanf " %d" Fn.id)

let g = DirectedGraph.create ~size:n
let () =
  Iter.(1 -- n) |> Iter.iter (fun i -> DirectedGraph.push g i a.(i - 1))

let self_cycle =
  Array.foldi a ~init:0 ~f:(fun i acc a -> acc + Bool.to_int (i + 1 = a))

let ans =
  DirectedGraph.strongly_connected_components g
  |> Iter.map Iter.length
  |> Iter.filter (fun len -> len > 1)
  |> Iter.sum

let ans = self_cycle + ans

let () = printf "%d\n" ans

val binary_search : 
  ?pos:int ->
  ?len:int ->
  'a t ->
  compare:('a -> 'key -> int) ->
  [ `Last_strictly_less_than
  | `Last_less_than_or_equal_to
  | `Last_equal_to
  | `First_equal_to
  | `First_greater_than_or_equal_to
  | `First_strictly_greater_than ] ->
  'key ->
  int option

let lower_bound, upper_bound =
  let arr_binsearch meth x = 
    Array.binary_search a ~compare meth x |> Option.value ~default:n
  in
  arr_binsearch `First_greater_than_or_equal_to,
  arr_binsearch `First_strictly_greater_than

let remove_range set l r =
  let compare = Int64.compare in
  let rec go acc set l =
    match Set.binary_search set ~compare `First_greater_than_or_equal_to l with
    | Some x when Int64.(x <= r) -> 
      go Set.(add acc x) (Set.remove set x) x
    | _ -> acc, set
  in
  go Int64.Set.empty set l

let x = [| (* ... *) |]
let n = Array.length x

let ans =
  Iter.(0 -- (1 lsl n - 1))
  |> Iter.map (fun bits ->
      Iter.(0 -- (n - 1))
      |> Iter.filter (fun i -> bits land (1 lsl i) <> 0)
      |> Iter.map (Array.get x)
      |> Iter.to_array
    )
  |> Iter.filter_map (function
      | [| (* ... *) |] -> Some ( (* ... *) )
      | _ -> None
    )

module Iter = struct
  include Iter
  let zip x y = flat_map (fun x -> map (fun y -> x, y) y) x
  let( * ) = zip
end

module type M = sig
  type t
  val compare : t -> t -> int
  val sexp_of_t : t -> Sexp.t
  val t_of_sexp : Sexp.t -> t
end
module Tuple2 = struct
  include Tuple2
  module Make (M1: M) (M2: M) = struct
    type t = M1.t * M2.t
    let compare = compare ~cmp1:M1.compare ~cmp2:M2.compare
    let sexp_of_t = sexp_of_t M1.sexp_of_t M2.sexp_of_t
    let t_of_sexp = t_of_sexp M1.t_of_sexp M2.t_of_sexp
  end
end
module S = Set.Make (Tuple2.Make (Int) (Int64))

let fib = Memo.recursive (module Int) @@ fun self -> function
  | n when n < 2 -> 1
  | n            -> self (n - 1) + self (n - 2)

let prime_count = Array.init 300_005 ~f:(const 0)
let rec sieve ?(acc=Iter.empty) = function
  | []      -> acc
  | x :: xs ->
    prime_count.(x) <- 1;
    List.filter ~f:(fun y -> y % x <> 0) xs
    |> sieve ~acc:(Iter.snoc acc @@ Int64.of_int x)
 
let primes = sieve (List.range 2 300_005) |> Iter.to_array
 
let () =
  for i = 1 to Array.length prime_count - 1 do
    prime_count.(i) <- prime_count.(i) + prime_count.(i - 1);
  done

let chmin i v = dp.(i) <- min dp.(i) v
let chmax i j v = dp.(i).(j) <- max dp.(i).(j) v

let ch f i v = dp.(i) <- f dp.(i) v

let ch f i j v = dp.(i).(j) <- f dp.(i).(j) v

module Array = struct
  include Array
  let ch f a i = a.(i) <- f a.(i)
  let incr = ch succ
end

let (.++()) a i = a.(i) <- a.(i) + 1
let (.--()) a i = a.(i) <- a.(i) - 1

module Heap = struct
  module X = struct
    type t = int * int * int
    let compare (a, _, _) (b, _, _) = Int.compare a b
  end

  type t = Leaf | Node of t * X.t * t * int

  let empty = Leaf

  let singleton k = Node (Leaf, k, Leaf, 1)
  let rank = function Leaf -> 0 | Node (_,_,_,r) -> r

  let rec merge t1 t2 =  
    match t1, t2 with
    | Leaf, t | t, Leaf -> t
    | Node (_, k1, _, _), Node (_, k2, _, _) 
      when 0 < X.compare k1 k2 -> merge t2 t1
    | Node (l, k, r, _), _ ->
      let r = merge r t2 in
      let rank_left = rank l and rank_right = rank r in
      let l, r, rank =
        if rank_left >= rank_right 
        then l, r, rank_right
        else r, l, rank_left
      in
      Node (l, k, r, rank + 1)

  let insert t x = merge (singleton x) t
  let find_min = function
    | Leaf              -> None
    | Node (_, k, _, _) -> Some k
  let del_min = function
    | Leaf              -> empty
    | Node (l, _, r, _) -> merge l r
end

let rec combinations ?(acc=Iter.empty) = function
  | 0 -> const @@ Iter.singleton Iter.empty
  | k -> generate acc k
and
  generate acc len iter = match Iter.head iter, Iter.drop 1 iter with
  | None, _     -> acc
  | Some hd, tl -> 
    let acc = combinations (len - 1) tl |> Iter.fold (fun acc cs -> Iter.snoc acc (Iter.cons hd cs)) acc in
    combinations ~acc len tl

let rec cartesian_power iter n =
  if n <= 0 then Iter.singleton []
  else
    Iter.flat_map
      (fun rest -> iter |> Iter.map (fun d -> d :: rest))
      (cartesian_power iter @@ n - 1)

let cumsum ~init ~f a =
  let paired f a b = let r = f a b in r, r in
  Array.append [| init |] @@ Array.folding_map a ~init ~f:(paired f)

let a = [|
  [| 1; 2; 3; 4; |];
  [| 5; 6; 7; 8; |];
|]

let cumsum ~init ~f a =
  let paired f a b = let r = f a b in r, r in
  Array.append [| init |] @@ Array.folding_map a ~init ~f:(paired f)

let cumsum2d ~init ~f a =
  let n = Array.length a and m = Array.length a.(0) in
  let a = Array.map a ~f:(cumsum ~init ~f) in
  for i = 1 to n - 1 do
    for j = 0 to m do
      a.(i).(j) <- f a.(i).(j) a.(i-1).(j)
    done
  done;
  Array.append [| Array.create ~len:(m+1) init |] a

let cumsum = cumsum2d ~init:0 ~f:Int.(+) a

let () =
  Array.iter cumsum ~f:(fun cumsum ->
    Array.iter cumsum ~f:(printf "%d ");
    printf "\n";
  )

(* output:
0 0 0 0 0
0 1 3 6 10
0 6 14 24 36
*) 

let solve (a, b) (c, d) =
  let a, b = a - 1, b - 1 in
  cumsum.(c).(d) - (cumsum.(a).(d) + cumsum.(c).(b)) + cumsum.(a).(b)

module ImosMake (M : sig include Int_intf.S end) = struct
  type t = M.t array

  let init n : t = Array.create ~len:(n + 2) M.zero

  let update t (x, y) =
    let l, r = min x y, max x y in
    t.(l) <- M.succ t.(l);
    t.(r) <- M.pred t.(r)

  let finish imos : t = 
    let cumsum ~init ~f a =
      let paired f a b = let r = f a b in r, r in
      Array.folding_map a ~init ~f:(paired f)
    in
    cumsum ~init:M.zero ~f:M.(+) imos

  let create n xs : t =
    let t = init n in
    List.iter xs ~f:(update t);
    finish t
end
module Imos = ImosMake(Int)

let imos =
  Imos.create len @@
  List.init n ~f:(fun _ ->
    scanf " %d %d" @@ fun a b -> a, b + 1
  )

module Imos2dMake (M : sig include Int_intf.S end) = struct
  type t = M.t array array

  let init n : t = Array.make_matrix ~dimx:(n+2) ~dimy:(n+2) M.zero

  let update (t:t) ((a, b), (c, d)) =
    let c, d = c + 1, d + 1 in
    t.(a).(b) <- M.succ t.(a).(b);
    t.(a).(d) <- M.pred t.(a).(d);
    t.(c).(b) <- M.pred t.(c).(b);
    t.(c).(d) <- M.succ t.(c).(d)

  let finish imos : t = 
    let cumsum ~init ~f a =
      let paired f a b = let r = f a b in r, r in
      Array.folding_map a ~init ~f:(paired f)
    in
    let cumsum2d ~init ~f a =
      let n = Array.length a and m = Array.length a.(0) in
      let a = Array.map a ~f:(cumsum ~init ~f) in
      for i = 1 to n - 1 do
        for j = 0 to m do
          a.(i).(j) <- f a.(i).(j) a.(i-1).(j)
        done
      done;
      Array.append [| Array.create ~len:(m+1) init |] a
    in
    cumsum2d ~init:M.zero ~f:M.(+) imos

  let create n xs : t =
    let t = init n in
    List.iter xs ~f:(update t);
    finish t

  let area t a b c d =
    let open M in
    t.(c).(d) - (t.(a).(d) + t.(c).(b)) + t.(a).(b)
end
module Imos2d = Imos2dMake(Int)

module ImosTriangleMake (M : sig include Int_intf.S end) = struct
  type t = M.t array array

  let init n : t = Array.make_matrix ~dimx:(n+5) ~dimy:(n+5) M.zero

  (* (a, b), (a + x, b), (a + x, b + x) を頂点とする正三角形 *)
  let update (t:t) (a, b, x) =
    t.(a).(b)         <- M.succ t.(a).(b);
    t.(a).(b+1)       <- M.pred t.(a).(b+1);
    t.(a+x+1).(b)     <- M.pred t.(a+x+1).(b);
    t.(a+x+1).(b+x+2) <- M.succ t.(a+x+1).(b+x+2);
    t.(a+x+2).(b+1)   <- M.succ t.(a+x+2).(b+1);
    t.(a+x+2).(b+x+2) <- M.pred t.(a+x+2).(b+x+2)

  let finish imos : t = 
    let cumsum ~init ~f a =
      let paired f a b = let r = f a b in r, r in
      Array.append [| init |] @@ Array.folding_map a ~init ~f:(paired f)
    in
    let cumsum2d ~init ~f a =
      let n = Array.length a and m = Array.length a.(0) in
      let a = Array.map a ~f:(cumsum ~init ~f) in
      for i = 1 to n - 1 do
        for j = 0 to m do
          a.(i).(j) <- f a.(i).(j) a.(i-1).(j)
        done
      done;
      for i = 1 to n - 1 do
        for j = 1 to m - 1 do
          a.(j).(i) <- f a.(j).(i) a.(j-1).(i-1)
        done
      done;
      Array.append [| Array.create ~len:(m+1) init |] a
    in
    cumsum2d ~init:M.zero ~f:M.(+) imos

  let create n xs : t =
    let t = init n in
    List.iter xs ~f:(update t);
    finish t

  let count_area =
    Array.fold ~init:0 ~f:(fun acc row ->
      acc + Array.count row ~f:M.(fun x -> zero < x)
    )
end
module ImosTriangle = ImosTriangleMake (Int)

let bfs m ~around ~f (init_k, init_v) =
  let tbl = Hashtbl.create m in
  Hashtbl.set tbl ~key:init_k ~data:init_v;
  let rec bfs q =
    match Fqueue.dequeue q with 
    | None        -> tbl
    | Some (v, q) ->
      let data = f @@ Hashtbl.find_exn tbl v in
      around v
      |> Iter.filter (fun   u -> Hashtbl.find tbl u |> Option.is_none)
      |> Iter.fold   (fun q u ->
        Hashtbl.set tbl ~key:u ~data;
        Fqueue.enqueue q u
      ) q
      |> bfs
  in
  bfs @@ Fqueue.singleton init_k