二叉搜索树 -> 数组 -> 二叉搜索树(Python/Java/C++/C/Go/JS/Rust)
由于输入的是一棵二叉搜索树,节点值满足 $左子树 < 根 < 右子树$,所以通过一次 94. 二叉树的中序遍历,把遍历到的节点值添加到一个数组中,可以直接得到一个递增数组,无需排序。
然后 108. 将有序数组转换为二叉搜索树,做法见 我的题解。
###py
class Solution:
# 94. 二叉树的中序遍历
def inorderTraversal(self, root: Optional[TreeNode]) -> List[int]:
def dfs(node: Optional[TreeNode]) -> None:
if node is None:
return
dfs(node.left) # 左
ans.append(node.val) # 根(这行代码移到前面就是前序,移到后面就是后序)
dfs(node.right) # 右
ans = []
dfs(root)
return ans
# 108. 将有序数组转换为二叉搜索树
def sortedArrayToBST(self, nums: List[int]) -> Optional[TreeNode]:
if not nums:
return None
m = len(nums) // 2
left = self.sortedArrayToBST(nums[:m])
right = self.sortedArrayToBST(nums[m + 1:])
return TreeNode(nums[m], left, right)
def balanceBST(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
nums = self.inorderTraversal(root)
return self.sortedArrayToBST(nums)
###py
class Solution:
# 94. 二叉树的中序遍历
def inorderTraversal(self, root: Optional[TreeNode]) -> List[int]:
def dfs(node: Optional[TreeNode]) -> None:
if node is None:
return
dfs(node.left) # 左
ans.append(node.val) # 根(这行代码移到前面就是前序,移到后面就是后序)
dfs(node.right) # 右
ans = []
dfs(root)
return ans
# 108. 将有序数组转换为二叉搜索树
def sortedArrayToBST(self, nums: List[int]) -> Optional[TreeNode]:
# 把 nums[left:right] 转成平衡二叉搜索树
def dfs(left: int, right: int) -> Optional[TreeNode]:
if left == right:
return None
m = (left + right) // 2
return TreeNode(nums[m], dfs(left, m), dfs(m + 1, right))
return dfs(0, len(nums))
def balanceBST(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
nums = self.inorderTraversal(root)
return self.sortedArrayToBST(nums)
###java
class Solution {
public TreeNode balanceBST(TreeNode root) {
List<Integer> nums = inorderTraversal(root);
return sortedArrayToBST(nums);
}
// 94. 二叉树的中序遍历
private List<Integer> inorderTraversal(TreeNode root) {
List<Integer> ans = new ArrayList<>();
dfs(ans, root);
return ans;
}
private void dfs(List<Integer> ans, TreeNode node) {
if (node == null) {
return;
}
dfs(ans, node.left); // 左
ans.add(node.val); // 根(这行代码移到前面就是前序,移到后面就是后序)
dfs(ans, node.right); // 右
}
// 108. 将有序数组转换为二叉搜索树
private TreeNode sortedArrayToBST(List<Integer> nums) {
return buildBST(nums, 0, nums.size());
}
// 把 nums[left] 到 nums[right-1] 转成平衡二叉搜索树
private TreeNode buildBST(List<Integer> nums, int left, int right) {
if (left == right) {
return null;
}
int m = (left + right) >>> 1;
return new TreeNode(nums.get(m), buildBST(nums, left, m), buildBST(nums, m + 1, right));
}
}
###cpp
class Solution {
// 94. 二叉树的中序遍历
vector<int> inorderTraversal(TreeNode* root) {
vector<int> ans;
// lambda 递归
auto dfs = [&](this auto&& dfs, TreeNode* node) -> void {
if (node == nullptr) {
return;
}
dfs(node->left); // 左
ans.push_back(node->val); // 根(这行代码移到前面就是前序,移到后面就是后序)
dfs(node->right); // 右
};
dfs(root);
return ans;
}
// 108. 将有序数组转换为二叉搜索树
TreeNode* sortedArrayToBST(vector<int>& nums) {
// 把 nums[left] 到 nums[right-1] 转成平衡二叉搜索树
auto dfs = [&](this auto&& dfs, int left, int right) -> TreeNode* {
if (left == right) {
return nullptr;
}
int m = left + (right - left) / 2;
return new TreeNode(nums[m], dfs(left, m), dfs(m + 1, right));
};
return dfs(0, nums.size());
}
public:
TreeNode* balanceBST(TreeNode* root) {
auto nums = inorderTraversal(root);
return sortedArrayToBST(nums);
}
};
###c
// 获取树的大小(节点个数)
int getSize(struct TreeNode* root) {
if (root == NULL) {
return 0;
}
return 1 + getSize(root->left) + getSize(root->right);
}
// 94. 二叉树的中序遍历
int* inorderTraversal(struct TreeNode* root, int* returnSize) {
int* ans = malloc(getSize(root) * sizeof(int));
*returnSize = 0;
void dfs(struct TreeNode* node) {
if (node == NULL) {
return;
}
dfs(node->left); // 左
ans[(*returnSize)++] = node->val; // 根(这行代码移到前面就是前序,移到后面就是后序)
dfs(node->right); // 右
}
dfs(root);
return ans;
}
// 108. 将有序数组转换为二叉搜索树
struct TreeNode* sortedArrayToBST(int* nums, int numsSize) {
// 把 nums[left] 到 nums[right-1] 转成平衡二叉搜索树
struct TreeNode* dfs(int left, int right) {
if (left == right) {
return NULL;
}
int m = left + (right - left) / 2;
struct TreeNode* node = malloc(sizeof(struct TreeNode));
node->val = nums[m];
node->left = dfs(left, m);
node->right = dfs(m + 1, right);
return node;
}
return dfs(0, numsSize);
}
struct TreeNode* balanceBST(struct TreeNode* root) {
int numsSize;
int* nums = inorderTraversal(root, &numsSize);
root = sortedArrayToBST(nums, numsSize);
free(nums);
return root;
}
###go
// 94. 二叉树的中序遍历
func inorderTraversal(root *TreeNode) (ans []int) {
var dfs func(*TreeNode)
dfs = func(node *TreeNode) {
if node == nil {
return
}
dfs(node.Left) // 左
ans = append(ans, node.Val) // 根(这行代码移到前面就是前序,移到后面就是后序)
dfs(node.Right) // 右
}
dfs(root)
return
}
// 108. 将有序数组转换为二叉搜索树
func sortedArrayToBST(nums []int) *TreeNode {
if len(nums) == 0 {
return nil
}
m := len(nums) / 2
return &TreeNode{
Val: nums[m],
Left: sortedArrayToBST(nums[:m]),
Right: sortedArrayToBST(nums[m+1:]),
}
}
func balanceBST(root *TreeNode) *TreeNode {
nums := inorderTraversal(root)
return sortedArrayToBST(nums)
}
###js
// 94. 二叉树的中序遍历
var inorderTraversal = function(root) {
function dfs(node) {
if (node === null) {
return;
}
dfs(node.left); // 左
ans.push(node.val); // 根(这行代码移到前面就是前序,移到后面就是后序)
dfs(node.right); // 右
}
const ans = [];
dfs(root);
return ans;
};
// 108. 将有序数组转换为二叉搜索树
var sortedArrayToBST = function(nums) {
// 把 nums[left] 到 nums[right-1] 转成平衡二叉搜索树
function dfs(left, right) {
if (left === right) {
return null;
}
const m = Math.floor((left + right) / 2);
return new TreeNode(nums[m], dfs(left, m), dfs(m + 1, right));
}
return dfs(0, nums.length);
};
var balanceBST = function(root) {
const nums = inorderTraversal(root)
return sortedArrayToBST(nums)
};
###rust
use std::rc::Rc;
use std::cell::RefCell;
impl Solution {
// 94. 二叉树的中序遍历
fn inorder_traversal(root: Option<Rc<RefCell<TreeNode>>>) -> Vec<i32> {
fn dfs(node: &Option<Rc<RefCell<TreeNode>>>, ans: &mut Vec<i32>) {
if let Some(node) = node {
let n = node.borrow();
dfs(&n.left, ans); // 左
ans.push(n.val); // 根(这行代码移到前面就是前序,移到后面就是后序)
dfs(&n.right, ans); // 右
}
}
let mut ans = vec![];
dfs(&root, &mut ans);
ans
}
// 108. 将有序数组转换为二叉搜索树
fn sorted_array_to_bst(nums: Vec<i32>) -> Option<Rc<RefCell<TreeNode>>> {
fn dfs(nums: &[i32]) -> Option<Rc<RefCell<TreeNode>>> {
if nums.is_empty() {
return None;
}
let m = nums.len() / 2;
Some(Rc::new(RefCell::new(TreeNode {
val: nums[m],
left: dfs(&nums[..m]),
right: dfs(&nums[m + 1..]),
})))
}
dfs(&nums)
}
pub fn balance_bst(root: Option<Rc<RefCell<TreeNode>>>) -> Option<Rc<RefCell<TreeNode>>> {
let nums = Self::inorder_traversal(root);
Self::sorted_array_to_bst(nums)
}
}
复杂度分析
- 时间复杂度:$\mathcal{O}(n)$,其中 $n$ 是二叉树的节点个数。注:Python 的第一种写法有切片的复制开销,二叉树的每一层都需要花费 $\mathcal{O}(n)$ 的时间,一共有 $\mathcal{O}(\log n)$ 层,所以时间复杂度是 $\mathcal{O}(n\log n)$;第二种写法避免了切片的复制开销,时间复杂度是 $\mathcal{O}(n)$。
- 空间复杂度:$\mathcal{O}(n)$。
专题训练
见下面树题单的「§2.9 二叉搜索树」和「§2.10 创建二叉树」。
分类题单
- 滑动窗口与双指针(定长/不定长/单序列/双序列/三指针/分组循环)
- 二分算法(二分答案/最小化最大值/最大化最小值/第K小)
- 单调栈(基础/矩形面积/贡献法/最小字典序)
- 网格图(DFS/BFS/综合应用)
- 位运算(基础/性质/拆位/试填/恒等式/思维)
- 图论算法(DFS/BFS/拓扑排序/基环树/最短路/最小生成树/网络流)
- 动态规划(入门/背包/划分/状态机/区间/状压/数位/数据结构优化/树形/博弈/概率期望)
- 常用数据结构(前缀和/差分/栈/队列/堆/字典树/并查集/树状数组/线段树)
- 数学算法(数论/组合/概率期望/博弈/计算几何/随机算法)
- 贪心与思维(基本贪心策略/反悔/区间/字典序/数学/思维/脑筋急转弯/构造)
- 链表、树与回溯(前后指针/快慢指针/DFS/BFS/直径/LCA)
- 字符串(KMP/Z函数/Manacher/字符串哈希/AC自动机/后缀数组/子序列自动机)
欢迎关注 B站@灵茶山艾府
