二叉搜索樹中第K小的元素[中等]

優質博文：IT-BLOG-CN

一、題目

給定一個二叉搜索樹的根節點root，和一個整數k，請你設計一個算法查找其中第k個最小元素（從1開始計數）。

示例 1：

輸入：root = [3,1,4,null,2], k = 1
輸出：1

示例 2：

輸入：root = [5,3,6,2,4,null,null,1], k = 3
輸出：3

樹中的節點數為n。
1 <= k <= n <= 104
0 <= Node.val <= 104

進階：如果二叉搜索樹經常被修改（插入/刪除操作）并且你需要頻繁地查找第 k 小的值，你將如何優化算法？

二、代碼

【1】中序遍歷： 二叉搜索樹具有如下性質：
? ● 結點的左子樹只包含小于當前結點的數。
? ● 結點的右子樹只包含大于當前結點的數。
? ● 所有左子樹和右子樹自身必須也是二叉搜索樹。

二叉樹的中序遍歷即按照訪問左子樹——根結點——右子樹的方式遍歷二叉樹；在訪問其左子樹和右子樹時，我們也按照同樣的方式遍歷；直到遍歷完整棵樹。

思路和算法： 因為二叉搜索樹和中序遍歷的性質，所以二叉搜索樹的中序遍歷是按照鍵增加的順序進行的。于是，我們可以通過中序遍歷找到第k個最小元素。

class Solution {public int kthSmallest(TreeNode root, int k) {Deque<TreeNode> stack = new ArrayDeque<TreeNode>();while (root != null || !stack.isEmpty()) {while (root != null) {stack.push(root);root = root.left;}root = stack.pop();--k;if (k == 0) {break;}root = root.right;}return root.val;}
}

時間復雜度： 時間復雜度：O(H+k)，其中H是樹的高度。在開始遍歷之前，我們需要O(H)到達葉結點。當樹是平衡樹時，時間復雜度取得最小值O(log?N+k)；當樹是線性樹（樹中每個結點都只有一個子結點或沒有子結點）時，時間復雜度取得最大值O(N+k)。
空間復雜度： O(H)，棧中最多需要存儲H個元素。當樹是平衡樹時，空間復雜度取得最小值O(log?N)；當樹是線性樹時，空間復雜度取得最大值O(N)。

【2】記錄子樹的結點數： 如果你需要頻繁地查找第k小的值，你將如何優化算法？

思路和算法： 在方法一中，我們之所以需要中序遍歷前k個元素，是因為我們不知道子樹的結點數量，不得不通過遍歷子樹的方式來獲知。因此，我們可以記錄下以每個結點為根結點的子樹的結點數，并在查找第k小的值時，使用如下方法搜索：
? ● 令node等于根結點，開始搜索。
? ● 對當前結點node進行如下操作：
? ? ○ 如果node的左子樹的結點數left小于k?1，則第k小的元素一定在node的右子樹中，令node等于其的右子結點，k等于k?left?1，并繼續搜索；
? ? ○ 如果node的左子樹的結點數left等于k?1，則第k小的元素即為node，結束搜索并返回node即可；
? ? ○ 如果node的左子樹的結點數left大于k?1，則第k小的元素一定在node的左子樹中，令node等于其左子結點，并繼續搜索。

在實現中，我們既可以將以每個結點為根結點的子樹的結點數存儲在結點中，也可以將其記錄在哈希表中。

class Solution {public int kthSmallest(TreeNode root, int k) {MyBst bst = new MyBst(root);return bst.kthSmallest(k);}
}class MyBst {TreeNode root;Map<TreeNode, Integer> nodeNum;public MyBst(TreeNode root) {this.root = root;this.nodeNum = new HashMap<TreeNode, Integer>();countNodeNum(root);}// 返回二叉搜索樹中第k小的元素public int kthSmallest(int k) {TreeNode node = root;while (node != null) {int left = getNodeNum(node.left);if (left < k - 1) {node = node.right;k -= left + 1;} else if (left == k - 1) {break;} else {node = node.left;}}return node.val;}// 統計以node為根結點的子樹的結點數private int countNodeNum(TreeNode node) {if (node == null) {return 0;}nodeNum.put(node, 1 + countNodeNum(node.left) + countNodeNum(node.right));return nodeNum.get(node);}// 獲取以node為根結點的子樹的結點數private int getNodeNum(TreeNode node) {return nodeNum.getOrDefault(node, 0);}
}

時間復雜度： 預處理的時間復雜度為O(N)，其中N是樹中結點的總數；我們需要遍歷樹中所有結點來統計以每個結點為根結點的子樹的結點數。搜索的時間復雜度為O(H)，其中H是樹的高度；當樹是平衡樹時，時間復雜度取得最小值O(log?N)；當樹是線性樹時，時間復雜度取得最大值O(N)。
空間復雜度： O(N)，用于存儲以每個結點為根結點的子樹的結點數。

【3】平衡二叉搜索樹： 如果二叉搜索樹經常被修改（插入/刪除操作）并且你需要頻繁地查找第k小的值，你將如何優化算法？

方法三需要先掌握平衡二叉搜索樹（AVL樹）的知識。平衡二叉搜索樹具有如下性質：
? ● 平衡二叉搜索樹中每個結點的左子樹和右子樹的高度最多相差1；
? ● 平衡二叉搜索樹的子樹也是平衡二叉搜索樹；
? ● 一棵存有 nnn 個結點的平衡二叉搜索樹的高度是O(log?n)。

思路和算法： 我們注意到在方法二中搜索二叉搜索樹的時間復雜度為O(H)，其中H是樹的高度；當樹是平衡樹時，時間復雜度取得最小值O(log?N)。因此，我們在記錄子樹的結點數的基礎上，將二叉搜索樹轉換為平衡二叉搜索樹，并在插入和刪除操作中維護它的平衡狀態。

class Solution {public int kthSmallest(TreeNode root, int k) {// 中序遍歷生成數值列表List<Integer> inorderList = new ArrayList<Integer>();inorder(root, inorderList);// 構造平衡二叉搜索樹AVL avl = new AVL(inorderList);// 模擬1000次插入和刪除操作int[] randomNums = new int[1000];Random random = new Random();for (int i = 0; i < 1000; ++i) {randomNums[i] = random.nextInt(10001);avl.insert(randomNums[i]);}shuffle(randomNums); // 列表亂序for (int i = 0; i < 1000; ++i) {avl.delete(randomNums[i]);}return avl.kthSmallest(k);}private void inorder(TreeNode node, List<Integer> inorderList) {if (node.left != null) {inorder(node.left, inorderList);}inorderList.add(node.val);if (node.right != null) {inorder(node.right, inorderList);}}private void shuffle(int[] arr) {Random random = new Random();int length = arr.length;for (int i = 0; i < length; i++) {int randIndex = random.nextInt(length);int temp = arr[i];arr[i] = arr[randIndex];arr[randIndex] = temp;}}
}// 平衡二叉搜索樹（AVL樹）：允許重復值
class AVL {Node root;// 平衡二叉搜索樹結點class Node {int val;Node parent;Node left;Node right;int size;int height;public Node(int val) {this(val, null);}public Node(int val, Node parent) {this(val, parent, null, null);}public Node(int val, Node parent, Node left, Node right) {this.val = val;this.parent = parent;this.left = left;this.right = right;this.height = 0; // 結點高度：以node為根節點的子樹的高度（高度定義：葉結點的高度是0）this.size = 1; // 結點元素數：以node為根節點的子樹的節點總數}}public AVL(List<Integer> vals) {if (vals != null) {this.root = build(vals, 0, vals.size() - 1, null);}}// 根據vals[l:r]構造平衡二叉搜索樹 -> 返回根結點private Node build(List<Integer> vals, int l, int r, Node parent) {int m = (l + r) >> 1;Node node = new Node(vals.get(m), parent);if (l <= m - 1) {node.left = build(vals, l, m - 1, node);}if (m + 1 <= r) {node.right = build(vals, m + 1, r, node);}recompute(node);return node;}// 返回二叉搜索樹中第k小的元素public int kthSmallest(int k) {Node node = root;while (node != null) {int left = getSize(node.left);if (left < k - 1) {node = node.right;k -= left + 1;} else if (left == k - 1) {break;} else {node = node.left;}}return node.val;}public void insert(int v) {if (root == null) {root = new Node(v);} else {// 計算新結點的添加位置Node node = subtreeSearch(root, v);boolean isAddLeft = v <= node.val; // 是否將新結點添加到node的左子結點if (node.val == v) { // 如果值為v的結點已存在if (node.left != null) { // 值為v的結點存在左子結點，則添加到其左子樹的最右側node = subtreeLast(node.left);isAddLeft = false;} else { // 值為v的結點不存在左子結點，則添加到其左子結點isAddLeft = true;}}// 添加新結點Node leaf = new Node(v, node);if (isAddLeft) {node.left = leaf;} else {node.right = leaf;}rebalance(leaf);}}// 刪除值為v的結點 -> 返回是否成功刪除結點public boolean delete(int v) {if (root == null) {return false;}Node node = subtreeSearch(root, v);if (node.val != v) { // 沒有找到需要刪除的結點return false;}// 處理當前結點既有左子樹也有右子樹的情況// 若左子樹比右子樹高度低，則將當前結點替換為右子樹最左側的結點，并移除右子樹最左側的結點// 若右子樹比左子樹高度低，則將當前結點替換為左子樹最右側的結點，并移除左子樹最右側的結點if (node.left != null && node.right != null) {Node replacement = null;if (node.left.height <= node.right.height) {replacement = subtreeFirst(node.right);} else {replacement = subtreeLast(node.left);}node.val = replacement.val;node = replacement;}Node parent = node.parent;delete(node);rebalance(parent);return true;}// 刪除結點p并用它的子結點代替它，結點p至多只能有1個子結點private void delete(Node node) {if (node.left != null && node.right != null) {return;// throw new Exception("Node has two children");}Node child = node.left != null ? node.left : node.right;if (child != null) {child.parent = node.parent;}if (node == root) {root = child;} else {Node parent = node.parent;if (node == parent.left) {parent.left = child;} else {parent.right = child;}}node.parent = node;}// 在以node為根結點的子樹中搜索值為v的結點，如果沒有值為v的結點，則返回值為v的結點應該在的位置的父結點private Node subtreeSearch(Node node, int v) {if (node.val < v && node.right != null) {return subtreeSearch(node.right, v);} else if (node.val > v && node.left != null) {return subtreeSearch(node.left, v);} else {return node;}}// 重新計算node結點的高度和元素數private void recompute(Node node) {node.height = 1 + Math.max(getHeight(node.left), getHeight(node.right));node.size = 1 + getSize(node.left) + getSize(node.right);}// 從node結點開始（含node結點）逐個向上重新平衡二叉樹，并更新結點高度和元素數private void rebalance(Node node) {while (node != null) {int oldHeight = node.height, oldSize = node.size;if (!isBalanced(node)) {node = restructure(tallGrandchild(node));recompute(node.left);recompute(node.right);}recompute(node);if (node.height == oldHeight && node.size == oldSize) {node = null; // 如果結點高度和元素數都沒有變化則不需要再繼續向上調整} else {node = node.parent;}}}// 判斷node結點是否平衡private boolean isBalanced(Node node) {return Math.abs(getHeight(node.left) - getHeight(node.right)) <= 1;}// 獲取node結點更高的子樹private Node tallChild(Node node) {if (getHeight(node.left) > getHeight(node.right)) {return node.left;} else {return node.right;}}// 獲取node結點更高的子樹中的更高的子樹private Node tallGrandchild(Node node) {Node child = tallChild(node);return tallChild(child);}// 重新連接父結點和子結點（子結點允許為空）private static void relink(Node parent, Node child, boolean isLeft) {if (isLeft) {parent.left = child;} else {parent.right = child;}if (child != null) {child.parent = parent;}}// 旋轉操作private void rotate(Node node) {Node parent = node.parent;Node grandparent = parent.parent;if (grandparent == null) {root = node;node.parent = null;} else {relink(grandparent, node, parent == grandparent.left);}if (node == parent.left) {relink(parent, node.right, true);relink(node, parent, false);} else {relink(parent, node.left, false);relink(node, parent, true);}}// trinode操作private Node restructure(Node node) {Node parent = node.parent;Node grandparent = parent.parent;if ((node == parent.right) == (parent == grandparent.right)) { // 處理需要一次旋轉的情況rotate(parent);return parent;} else { // 處理需要兩次旋轉的情況：第1次旋轉后即成為需要一次旋轉的情況rotate(node);rotate(node);return node;}}// 返回以node為根結點的子樹的第1個元素private static Node subtreeFirst(Node node) {while (node.left != null) {node = node.left;}return node;}// 返回以node為根結點的子樹的最后1個元素private static Node subtreeLast(Node node) {while (node.right != null) {node = node.right;}return node;}// 獲取以node為根結點的子樹的高度private static int getHeight(Node node) {return node != null ? node.height : 0;}// 獲取以node為根結點的子樹的結點數private static int getSize(Node node) {return node != null ? node.size : 0;}
}