treeMap

完整代码

---

class Node {
  left;
  right;
  parent;
  height;
  constructor(key, val) {
    this.key = key;
    this.val = val;
    this.height = 0;
  }
  compareTo(target) {
    let targetKey = target;
    if (target instanceof Node) {
      targetKey = target.key;
    }
    return  this.key - targetKey.key;
  }
}
class TreeSet {
  constructor() {
    this.size = 0;
    this.root = null;
  }
  balanceFactor(node) {
    if (node == null) {
      return 0;
    }
    return this.height(node.left) - this.height(node.right);
  }
  height(node) {
    if (node == null) {
      return 0;
    }
    return node.height;
  }
  updateHeight(node) {
    node.height = Math.max(this.height(node.left), this.height(node.right)) + 1;
  }
  leftRotate(node) {
    let right = node.right;
    node.right = right.left;
    if (right.left) {
      right.left.parent = node;
    }
    right.parent = node.parent;
    if (!node.parent) {
      this.root = right;
    } else if (node.parent.left == node) {
      node.parent.left = right;
    } else {
      node.parent.right = right;
    }
    right.left = node;
    node.parent = right;
    this.updateHeight(node);
    this.updateHeight(right);
  }
  rightRotate(node) {
    let left = node.left;
    node.left = left.right;
    if (left.right) {
      left.right.parent = node;
    }
    left.parent = node.parent;
    if (!node.parent) {
      this.root = left;
    } else if (node.parent.left == node) {
      node.parent.left = left;
    } else {
      node.parent.right = left;
    }
    left.right = node;
    node.parent = left;
    this.updateHeight(node);
    this.updateHeight(left);
  }
  isBalanced(node) {
    if (!node) {
      return true;
    }
    let balanceFactor = this.balanceFactor(node);
    if (Math.abs(balanceFactor) > 1) {
      return false;
    }
    return this.isBalanced(node.left) && this.isBalanced(node.right);
  }
  balance(node) {
    let current = node;
    while (current) {
      this.updateHeight(current);
      if (this.balanceFactor(current) > 1) {
        if (this.balanceFactor(current.left) < 0) {
          this.leftRotate(current.left);
        }
        this.rightRotate(current);
      } else if (this.balanceFactor(current) < -1) {
        if (this.balanceFactor(current.right) > 0) {
          this.rightRotate(current.right);
        }
        this.leftRotate(current);
      } else {
        current = current.parent;
      }
    }
  }
  transplant(aNode, bNode) {
    if (!aNode.parent) {
      this.root = bNode;
    } else if (aNode.parent.left === aNode) {
      aNode.parent.left = bNode;
    } else {
      aNode.parent.right = bNode;
    }
    if (bNode) {
      bNode.parent = aNode.parent;
    }
  }
  add(key, value) {
    let parent = null;
    let current = this.root;
    this.size++;
    while (current) {
      parent = current;
      if (current.compareTo(key) > 0) {
        current = current.left;
      } else if (current.compareTo(key) < 0) {
        current = current.right;
      } else {
        current.key = key;
        current.value = value;
        this.size--;
        return;
      }
    }
    const node = new Node(key, value);
    if (!parent) {
      this.root = node;
    } else {
      if (node.compareTo(parent) < 0) {
        parent.left = node;
      } else {
        parent.right = node;
      }
      node.parent = parent;
    }
    this.balance(node);
  }
  remove(key) {
    let node = this.getNode(key);
    if (node) {
      return this.removeNode(node);
    }
    return false;
  }
  removeNode(node) {
    let remove = node;
    let replace;
    if (!node.left) {
      replace = node.right;
      this.transplant(node, node.right);
    } else if (!node.right) {
      replace = node.left;
      this.transplant(node, node.left);
    } else {
      remove = this.getSucceed(node);
      replace = remove.right;
      if (remove.parent == node) {
        if (replace) {
          replace.parent = remove;
        }
      } else {
        this.transplant(remove, remove.right);
        remove.right = node.right;
        remove.right.parent = remove;
      }
      this.transplant(node, remove);
      remove.left = node.left;
      remove.left.parent = remove;
    }
    this.balance(replace);
    this.size--;
    return true;
  }
  get(key) {
    const node = this.getNode(key);
    return node;
  }
  first() {
    let current = this.root;
    while (current.left) {
      current = current.left;
    }
    return current;
  }
  getNode(key) {
    let current = this.root;
    while (current) {
      if (current.compareTo(key) > 0) {
        current = current.left;
      } else if (current.compareTo(key) < 0) {
        current = current.right;
      } else {
        return current;
      }
    }
    return current;
  }
  getSucceed(node) {
    let right = node.right;
    let succeed = right;
    while (right) {
      succeed = right;
      right = right.left;
    }
    return succeed;
  }
  preOrderTraverse() {
    const result = [];
    if (!this.root) {
      return result;
    }
    const stack = [];
    let currentNode = this.root;
    while (currentNode || stack.length) {
      while (currentNode) {
        stack.push(currentNode);
        result.push(currentNode.key);
        currentNode = currentNode.left;
      }
      currentNode = stack.pop();
      currentNode = currentNode.right;
    }
    return result;
  }
}

测试用例

---

const treeSet = new TreeSet();
// const array = [1, 2];
const array = [1, 2, 3];
array.forEach((item) => {
  treeSet.add(item);
});
console.log(treeSet.preOrderTraverse());
// treeSet.remove(1);
// treeSet.remove(2);
// console.log(treeSet.preOrderTraverse());

console.log(treeSet.first());

参考资料

---

java实现AVL树

AVL树(三)之 Java的实现