669. 修剪二叉搜索树
问题描述:传送门
思路:
比较无脑的想法是:
递归处理,遇到root->val < low || root->val >high 的时候,直接return NULL。
如下:
class Solution {
public:
TreeNode* trimBST(TreeNode* root, int low, int high) {
if (root == nullptr || root->val < low || root->val > high)
return nullptr;
root->left = trimBST(root->left, low, high);
root->right = trimBST(root->right, low, high);
return root;
}
};
但仔细分析题目,【1,3】区间在二叉搜索树中不是节点0和节点3就能进行剪枝操作,还要考虑节点0的右子树。
如下图:
从图里也能看出来需要重构二叉树。
比如上图的0不符合要求,那就将节点0的右孩子节点2给节点3的左孩子。
完成节点0的删除,如下图:
1、递归
① 确定递归函数的参数和返回值
参数:根节点,最低限,最高限。
返回值:返回值移除节点。
代码如下:
TreeNode* trimBST(TreeNode* root, int low, int high)
② 确定终止条件
遇到空节点返回,因为修剪操作不是在终止条件上进行的。
if (root == nullptr ) return nullptr;
③ 确定单层递归逻辑
如果root(当前节点)的元素小于low的数值,那么应该递归右子树,并返回右子树中,符合条件的头结点。
if (root->val < low) {
TreeNode* right = trimBST(root->right, low, high);
// 寻找符合区间[low, high]的节点
return right;
}
如果root(当前节点)的元素大于high的,那么应该递归左子树,并返回左子树中,符合条件的头结点。
if (root->val > high) {
TreeNode* left = trimBST(root->left, low, high);
// 寻找符合区间[low, high]的节点
return left;
}
将下一层处理完左子树的结果赋给root->left,处理完右子树的结果赋给root->right。
最后返回root节点,代码如下:
root->left = trimBST(root->left, low, high);
// root->left接入符合条件的左孩子
root->right = trimBST(root->right, low, high);
// root->right接入符合条件的右孩子
return root;
再回顾刚才我画的图,
如下代码就相当于把节点0的右孩子(节点2)返回给上一层,
if (root->val < low) {
TreeNode* right = trimBST(root->right, low, high);
// 寻找符合区间[low, high]的节点
return right;
}
然后如下代码相当于用节点3的左孩子,把下一层返回的 节点0的右孩子(节点2)接住。
root->left = trimBST(root->left, low, high);
此时,节点3的左孩子就变成了节点2,将节点0从二叉树中移除了。
最后,整体代码如下:
class Solution {
public:
TreeNode* trimBST(TreeNode* root, int low, int high) {
if (root == nullptr ) return nullptr;
if (root->val < low) {
TreeNode* right = trimBST(root->right, low, high); // 寻找符合区间[low, high]的节点
return right;
}
if (root->val > high) {
TreeNode* left = trimBST(root->left, low, high); // 寻找符合区间[low, high]的节点
return left;
}
root->left = trimBST(root->left, low, high); // root->left接入符合条件的左孩子
root->right = trimBST(root->right, low, high); // root->right接入符合条件的右孩子
return root;
}
};
精简后的代码如下:
class Solution {
public:
TreeNode* trimBST(TreeNode* root, int low, int high) {
if (root == nullptr) return nullptr;
if (root->val < low) return trimBST(root->right, low, high);
if (root->val > high) return trimBST(root->left, low, high);
root->left = trimBST(root->left, low, high);
root->right = trimBST(root->right, low, high);
return root;
}
};
2、迭代
在剪枝的时候,分三步:
- 将root移动到【L,R】范围内,注意左闭右闭区间。
- 剪枝左子树
- 剪枝右子树
代码如下:
class Solution {
public:
TreeNode* trimBST(TreeNode* root, int L, int R) {
if (!root) return nullptr;
// 处理头结点,让root移动到[L, R] 范围内,注意是左闭右闭
while (root->val < L || root->val > R) {
if (root->val < L) root = root->right; // 小于L往右走
else root = root->left; // 大于R往左走
}
TreeNode *cur = root;
// 此时root已经在[L, R] 范围内,处理左孩子元素小于L的情况
while (cur != nullptr) {
while (cur->left && cur->left->val < L) {
cur->left = cur->left->right;
}
cur = cur->left;
}
cur = root;
// 此时root已经在[L, R] 范围内,处理右孩子大于R的情况
while (cur != nullptr) {
while (cur->right && cur->right->val > R) {
cur->right = cur->right->left;
}
cur = cur->right;
}
return root;
}
};