資料結構

1class AVLNode {
2  left: AVLNode | null = null;
3  right: AVLNode | null = null;
4  height = 1;
5  constructor(public key: number) {}
6}
7 
8const height = (n: AVLNode | null) => (n ? n.height : 0);
9const balanceOf = (n: AVLNode) => height(n.left) - height(n.right);
10 
11function update(n: AVLNode): void {
12  n.height = 1 + Math.max(height(n.left), height(n.right));
13}
14 
15function rotateRight(y: AVLNode): AVLNode {
16  const x = y.left!;
17  y.left = x.right;
18  x.right = y;
19  update(y);
20  update(x);
21  return x;
22}
23 
24function rotateLeft(x: AVLNode): AVLNode {
25  const y = x.right!;
26  x.right = y.left;
27  y.left = x;
28  update(x);
29  update(y);
30  return y;
31}
32 
33function rebalance(n: AVLNode): AVLNode {
34  update(n);
35  const b = balanceOf(n);
36  if (b > 1) {
37    if (balanceOf(n.left!) < 0) n.left = rotateLeft(n.left!);
38    return rotateRight(n);
39  }
40  if (b < -1) {
41    if (balanceOf(n.right!) > 0) n.right = rotateRight(n.right!);
42    return rotateLeft(n);
43  }
44  return n;
45}
46 
47function insert(node: AVLNode | null, key: number): AVLNode {
48  if (!node) return new AVLNode(key);
49  if (key === node.key) return node;
50  if (key < node.key) node.left = insert(node.left, key);
51  else node.right = insert(node.right, key);
52  return rebalance(node);
53}
54 
55function remove(node: AVLNode | null, key: number): AVLNode | null {
56  if (!node) return null;
57  if (key < node.key) node.left = remove(node.left, key);
58  else if (key > node.key) node.right = remove(node.right, key);
59  else if (!node.left || !node.right) return node.left ?? node.right;
60  else {
61    let min = node.right;
62    while (min.left) min = min.left;
63    node.key = min.key;
64    node.right = remove(node.right, min.key);
65  }
66  return rebalance(node);
67}