[Leetcode 450] Delete Node in a BST

作者 Shilei Tian 日期 2017-02-15
[Leetcode 450] Delete Node in a BST

题目要求

Given a root node reference of a BST and a key, delete the node with the given key in the BST. Return the root node reference (possibly updated) of the BST.

Basically, the deletion can be divided into two stages:

  1. Search for a node to remove.
  2. If the node is found, delete the node.

Note: Time complexity should be O(height of tree).

Example:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
root = [5,3,6,2,4,null,7]
key = 3

5
/ \
3 6
/ \ \
2 4 7

Given key to delete is 3. So we find the node with value 3 and delete it.

One valid answer is [5,4,6,2,null,null,7], shown in the following BST.

5
/ \
4 6
/ \
2 7

Another valid answer is [5,2,6,null,4,null,7].

5
/ \
2 6
\ \
4 7

解题思路

其实这道题目本身没什么好讲的,删除 BST 中的结点不涉及翻转,所以相对简单,但是想做地高效一点,还是需要注意编程的。我写出来的第一个版本运行时间是 65ms,紧接着只修改了其中的一个函数,就得到 39ms 的运行时间,然后将原先的算法改成递归来做了,运行时间变成 36ms,已经超过了 84.81% 的人了,虽然还没有上第一梯度,仅有 32ms

第一个版本的代码如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
class Solution {
public:
TreeNode* deleteNode(TreeNode* root, int key) {
TreeNode* target = nullptr;
if (findNode(root, key, &target)) {
if (target->left != nullptr && target->right != nullptr) {
auto willCut = findMin(target->right);
target->val = willCut->val;
cut(root, willCut);
} else {
auto parent = findParent(root, target);
if (parent == nullptr) {
return target->left == nullptr ? target->right : target->left;
}
auto next = (target->left == nullptr ? target->right : target->left);
if (parent->left == target) {
parent->left = next;
} else {
parent->right = next;
}
}
}
return root;
}

private:
bool findNode(TreeNode* root, const int key, TreeNode** res) {
if (root == nullptr) {
return false;
} else if (root->val == key) {
*res = root;
return true;
}
return findNode(root->left, key, res) || findNode(root->right, key, res);
}

TreeNode* findParent(TreeNode *root, TreeNode *target) {
if (root == nullptr) {
return nullptr;
} else if (root->left == target || root->right == target) {
return root;
}
auto res = findParent(root->left, target);
if (res != nullptr) {
return res;
}
res = findParent(root->right, target);
return res;
}

TreeNode* findMin(TreeNode* root) {
if (root->left) {
return findMin(root->left);
} else {
return root;
}
}

void cut(TreeNode* root, TreeNode* target) {
auto parent = findParent(root, target);
if (parent->right == target) {
parent->right = target->right;
} else {
parent->left = target->right;
}
}
};

整体的思路就是先找到这个结点,然后再判断它是否是有两个孩子,如果是的话,就从右子树中找到最小的那个结点,用它的值去替换那个要删除结点的值,然后再删除这个最小的结点。

代码很明确,分成几个函数,运行时间很差,69ms,因此我突然想到那个 findMin 函数可以改成非递归的,因为这个只检查左子树是否存在。所以就有了下面第二个版本(这里只贴新的 findMin 函数,因为只改了这一点):

1
2
3
4
5
6
TreeNode* findMin(TreeNode* root) {
while (root->left != nullptr) {
root = root->left;
}
return root;
}

消耗的时间却是到了 36ms,这个确实令我有些惊讶。我分析可能是由于这个函数被调用很多次吧,所以递归的方法在这里耗费了太长时间。

然后我又看了一下 Solutions 里面大家的解法,排名第一的是一个递归的方法。曾经一名微软的人跟我说,二叉树的问题,“十个有十个是用递归做”。当然,这句话强调了递归思想在二叉树题目中的地位。因此就有了下面这个版本:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class Solution {
public:
TreeNode* deleteNode(TreeNode* root, int key) {
if (root == nullptr) {
return root;
} else if (key > root->val) {
root->right = deleteNode(root->right, key);
} else if (key < root->val) {
root->left = deleteNode(root->left, key);
} else {
if (root->left == nullptr) {
return root->right;
} else if (root->right == nullptr) {
return root->left;
}
auto minNode = findMin(root->right);
root->val = minNode->val;
root->right = deleteNode(root->right, root->val);
}
return root;
}

private:

TreeNode* findMin(TreeNode* root) {
while (root->left != nullptr) {
root = root->left;
}
return root;
}
};

代码很简短,只有第一个版本的一半不到,运行时间可以到 36ms,又提升了一点。