diff --git a/konnovaea/phonebook.py b/konnovaea/phonebook.py index d38d076..169190e 100644 --- a/konnovaea/phonebook.py +++ b/konnovaea/phonebook.py @@ -120,22 +120,29 @@ def _bst_find_mine(node): return current def bst_delete(root, name): + if root is None: return None + if name < root['name']: root['left'] = bst_delete(root['left'], name) elif name > root['name']: root['right'] = bst_delete(root['right'], name) - else: - if root['left'] is None: - return root['right'] - elif root['right'] is None: - return root['left'] + return root + + if root['left'] is None: + return root['right'] + elif root['right'] is None: + return root['left'] - successor = _bst_find_mine(root['right']) - root['name'] = successor['name'] - root['phone'] = successor['phone'] - root['right'] = bst_delete(root['right'], successor['name']) + current = root['right'] + while current['left'] is not None: + current = current['left'] + + root['name'] = current['name'] + root['phone'] = current['phone'] + + root['right'] = bst_delete(root['right'], current['name']) return root