確か、いい感じに深くなりすぎない木だよね。回転?なにそれ?みたいな状態から実装し始めたら、面倒で死ぬかと思った。
AVL Tree by Java -- これで分かったAVL木 とか
AVL木の回転(要素の挿入や削除のしかた) - Qiita とかを参考に実装した。2週間くらい前から2回ほどチャレンジしてギブアップしていて、本日、腰を据えて三度目の正直。コードを整形する気力が残ってないので、そのまま貼り付け。
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'
;;
実装してて思ったんだけど、左回転して右回転・・・のような話は、どこかで聞いたことがある気がする。大昔になんかの本で読んだのだろうか。