One approach that uses O(n) extra space, is to store for each node N, the pointer to the nodes with minimum and maximum values in the sub-tree rooted at N. Let's denote the minimum node rooted at N by N.min and the maximum node by N.max. Then for a sub-tree rooted at N, the sub-tree has a "defect", if :
N.val < N.left.max.val and/or N.val > N.right.min.val
This is quite intuitive as either of the above condition violates the binary search tree principle.
The maximum value on the left sub-tree should be less than the root and the minimum value on the right sub-tree should be greater than the root.
In the above BST, we swap the nodes 20 and 21. Node that for the node 21, the condition :
N.val > N.right.min.val
holds true, because N.val = 21 and N.right.min.val = 20. But the condition :
N.val < N.left.max.val
do not hold true, because N.val = 21 and N.left.max.val = 18.
Following is the code for getting the min-max values. We get the min-max values as a map because we will assume that the class for the TreeNode is defined elsewhere and it is immutable.
def get_min_max(self, root, min_max_map): if root.left is not None: min_left, max_left = self.get_min_max(root.left, min_max_map) else: min_left, max_left = None, None if root.right is not None: min_right, max_right = self.get_min_max(root.right, min_max_map) else: min_right, max_right = None, None a = root if min_left is not None and a.val > min_left.val: a = min_left if min_right is not None and a.val > min_right.val: a = min_right b = root if max_left is not None and b.val < max_left.val: b = max_left if max_right is not None and b.val < max_right.val: b = max_right min_max_map[root] = (a, b) return a, b
If either of the condition holds true, then we swap the values of root N of the sub-tree with N.left.max or N.right.min depending on which condition holds true, else if both the condition holds, then swap N.left.max with N.right.min.
In the above BST, the nodes 13 and 23 are swapped from the original. For the root node 15, note that both the conditions holds true, because 15 > 13 and 15 < 23. In such cases, we swap the nodes with values 13 and 23 and do not touch the root node.
The code for adjustment of the nodes is done by the following code:
def swap(self, node1, node2): if node1 is not None and node2 is not None: temp = node1.val node1.val = node2.val node2.val = temp def recover(self, root, min_max_map): p, q = min_max_map[root.left], min_max_map[root.right] x, y = root.val < p.val, root.val > q.val if x and y: self.swap(p, q) elif x: self.swap(root, p) elif y: self.swap(root, q) else: self.recover(root.left, min_max_map) self.recover(root.right, min_max_map)
In order to achieve the same result as above but without using any extra space, i.e. in O(1) space, we need to combine the 'get_min_max' and 'recover' steps into a single function.
There would be no change in the logic of the code, just that, in the 'get_min_max' function, we need to additionally track the sub-trees where the defects occurred, so that we don't need to traverse again to get the sub-trees where the defects occurred (in 'recover' method).
Below is the modified 'get_min_max' method :
def get_min_max(self, root): has_defect, last_defect = 0, (None, None) if root.left is not None: min_left, max_left, has_defect_left, last_defect_left = self.get_min_max(root.left) else: min_left, max_left, has_defect_left, last_defect_left = None, None, 0, (None, None) if root.right is not None: min_right, max_right, has_defect_right, last_defect_right = self.get_min_max(root.right) else: min_right, max_right, has_defect_right, last_defect_right = None, None, 0, (None, None) if max_left is not None and max_left.val > root.val: last_defect_left = (root, max_left) has_defect_left = 1 has_defect, last_defect = 1, last_defect_left if min_right is not None and min_right.val < root.val: last_defect_right = (root, min_right) has_defect_right = 1 has_defect, last_defect = 1, last_defect_right if has_defect_left == 1 and has_defect_right == 1: self.swap(last_defect_left, last_defect_right) has_defect, last_defect = 0, (None, None) elif has_defect_left == 1: has_defect, last_defect = 1, last_defect_left elif has_defect_right == 1: has_defect, last_defect = 1, last_defect_right a = root if min_left is not None and a.val > min_left.val: a = min_left if min_right is not None and a.val > min_right.val: a = min_right b = root if max_left is not None and b.val < max_left.val: b = max_left if max_right is not None and b.val < max_right.val: b = max_right return a, b, has_defect, last_defect def recoverTree(self, root): a, b, has_defect, last_defect = self.get_min_max(root) if has_defect == 1: self.swap(last_defect, last_defect)
We track the sub-tree(s) where the defect occurred using the 'last_defect' variable.
Observe that, if a tree T with root N, has a defect within it, then there could at-least one other sub-tree inside it which will have a defect. While recursing bottom-up, we should not swap nodes immediately on observing a defect because, there might be a parent node which needs to swapped with the defective node.
When we swap the nodes 2 with node 4, note that there is a defect in the sub-tree rooted at 3 as well as the sub-tree rooted at 2. Thus during the bottom up recursion, if we swap 3 with 4, then it will be incorrect since it's super-parent node 2 also has a defect, which can give the correct fix.
We can't swap any node values till we have seen all the nodes all the way up to the root node.
Only in the case where we observe that both the left and the right sub-trees contains a defect, then we can safely swap these nodes immediately on discovering them. Both the sub-trees rooted at 20 has a defect and thus we can safely swap them instead of deferring the swapping till we have seen the root node.
Run time complexity of the above code is O(n) and space complexity is O(1), because we are not tracking the min-max values separately in memory.
Categories: PROBLEM SOLVING