AVL木

確か、いい感じに深くなりすぎない木だよね。回転？なにそれ？みたいな状態から実装し始めたら、面倒で死ぬかと思った。

AVL Tree by Java -- これで分かったAVL木 とか AVL木の回転(要素の挿入や削除のしかた) - Qiita とかを参考に実装した。２週間くらい前から２回ほどチャレンジしてギブアップしていて、本日、腰を据えて三度目の正直。コードを整形する気力が残ってないので、そのまま貼り付け。

type 'a btree =
| Node of int * 'a * 'a btree * 'a btree
| Leaf
;;

let rotate_right t =
  match t with
  | Leaf -> Leaf
  | Node (bx, x, Node (by, y, ty1, ty2), tx2) ->
      let hty1 = by in
      let hty2 = 0 in
      let hy = max hty1 hty2 + 1 in
      let htx2 = hy - bx in
      let bx' = hty2 - htx2 in
      let hx' = max hty2 htx2 + 1 in
      let by' = hty1 - hx' in
      Node(by', y, ty1, Node(bx', x, ty2, tx2))
  | _ -> failwith "can't rotate right"
;;

let rotate_left t =
  match t with
  | Leaf -> Leaf
  | Node (bx, x, tx1, Node (by, y, ty1, ty2)) ->
      let hty1 = by in
      let hty2 = 0 in
      let hy = max hty1 hty2 + 1 in
      let htx1 = hy + bx in
      let bx' = htx1 - hty1 in
      let hx' = max htx1 hty1 + 1 in
      let by' = hx' - hty2 in
      Node(by', y, Node(bx', x, tx1, ty1), ty2)
  | _ -> failwith "can't rotate left"
;;

let rotate_left_right t =
  match t with
  | Node (bx, x, Node(by, y, ty1, ty2), tx2) ->
      rotate_right (Node (bx, x, Node (by, y, ty1, rotate_left ty2), tx2))
  | _ -> failwith "can't rotate left right"
;;

let rotate_right_left t =
  match t with
  | Node (bx, x, tx1, Node(by, y, ty1, ty2)) ->
      rotate_left (Node (bx, x, tx1, Node (by, y, rotate_right ty1, ty2)))
  | _ -> failwith "can'trotate right left"
;;

let bias t =
  match t with
  | Leaf -> 0
  | Node (bx, _, _, _) -> bx

let insert_avl t a =
  let rec insert_avl' t a =
    match t with
    | Leaf -> (Node (0, a, Leaf, Leaf), true)
    | Node (bx, x, tx1, tx2) ->
        if a < x then let (tx1', is_grown) = insert_avl' tx1 a in
                      let bx' = bx + if is_grown then 1 else 0 in
                      let t' = Node (bx', x, tx1', tx2) in
                      if is_grown then if bx' == 2 then if bias tx2 == -1 then (rotate_left_right t', false)
                                                                          else (rotate_right t', false)
                                                   else (t', bx' != 0)
                                  else (t', false)
                 else let (tx2', is_grown) = insert_avl' tx2 a in
                      let bx' = bx - if is_grown then 1 else 0 in
                      let t' = Node (bx', x, tx1, tx2') in
                      if is_grown then if bx' == -2 then if bias tx1 == 1 then (rotate_right_left t', false)
                                                                          else (rotate_left t', false)
                                                    else (t', bx' != 0)
                                  else (t', false)
  in let (t', _) = insert_avl' t a in t'
;;

実装してて思ったんだけど、左回転して右回転・・・のような話は、どこかで聞いたことがある気がする。大昔になんかの本で読んだのだろうか。