セグ木

segtree.ml

let rec fold_tournament dir f = function
  | [x] -> x
  | xs ->
      List.fold_left (fun (acc, prev) x ->
        match prev with
        | None -> (acc, Some x)
        | Some y -> ((if dir then f y x else f x y) :: acc, None)) ([], None) xs
      |> (function
          | (acc, None) -> acc
          | (acc, Some x) -> x :: acc)
      |> fold_tournament (not dir) f
(* fold_tournament ( * ) [x1; x2; x3; x4; x5; x6; x7; x8 ... ] = (... (((x1 * x2) * (x3 * x4)) * ((x5 * x6) * (x7 * x8))) ...) * ( ... ) *)
let fold_tournament f xs = fold_tournament true f xs

module type SemiGroup = sig
  type t
  val op : t -> t -> t
end

module Make (S : SemiGroup) : sig
  type t
  type elt
  (* リストからセグ木を作る *)
  val of_list : elt list -> t
  (*
   * update i x t
   * i番目の要素をxに変更したセグ木を作る
   *)
  val update : int -> elt -> t -> t
  (*
   * find l r t
   * [l, r)の要素の積を求める
   *)
  val find : int -> int -> t -> elt
end with type elt = S.t = struct
  type elt = S.t
  type t =
    | Leaf of elt
    | Node of { size : int; product : elt; left : t ; right : t }

  (* セグ木の要素数 *)
  let size = function
    | Leaf _ -> 1
    | Node { size } -> size

  (* セグ木の持っている要素全ての積 *)
  let product = function
    | Leaf product
    | Node { product } -> product

  let of_list l =
    fold_tournament (fun left right ->
      Node
        { size = size left + size right;
          product = S.op (product left) (product right);
          left; right }) (List.map (fun x -> Leaf x) l)

  let rec update i x t =
    assert (0 <= i && i < size t);
    match t with
    | Leaf _ -> Leaf x
    | Node ({ left; right } as record) ->
        if i < size left then
          let left' = update i x left in
          Node { record with
            product = S.op (product left') (product right);
            left = left' }
        else
          let right' = update (i - size left) x right in
          Node { record with
            product = S.op (product left) (product right');
            right = right' }

  let rec find l r t =
    assert (0 <= l && l < r && r <= size t);
    match t with
    | Leaf x -> x
    | Node { left; right } ->
        if l = 0 && r = size t - 1 then product t
        else if r <= size left then find l r left
        else if size left <= l then find (l - size left) (r - size left) right
        else S.op (find l (size left) left) (find 0 (r - size left) right)
end