Last active
December 17, 2021 20:53
-
-
Save davidinga/a8f1d2903cfec2326982c5835d959d30 to your computer and use it in GitHub Desktop.
Simple BST in Swift
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
class Program { | |
class BST { | |
var value: Int | |
var left: BST? | |
var right: BST? | |
init(value: Int) { | |
self.value = value | |
left = nil | |
right = nil | |
} | |
func insert(value: Int) -> BST { | |
var node: BST? = self | |
while node != nil { | |
if value < node!.value, node?.left == nil { | |
node!.left = BST(value: value) | |
break | |
} | |
if value >= node!.value, node?.right == nil { | |
node!.right = BST(value: value) | |
break | |
} | |
if value < node!.value { | |
node = node?.left | |
} else { | |
node = node?.right | |
} | |
} | |
return self | |
} | |
// Average: O(log(n)) time | O(1) space | |
// Worst: O(n) time | O(1) space | |
func contains(value: Int) -> Bool { | |
var node: BST? = self | |
while node != nil { | |
if node?.value == value { return true } | |
if value < node!.value { | |
node = node?.left | |
} else { | |
node = node?.right | |
} | |
} | |
return false | |
} | |
func remove(value: Int?, parentNode: BST?) -> BST { | |
guard let value = value else { return self } | |
var currentNode: BST? = self | |
var parentNode = parentNode | |
while let node = currentNode { | |
if value < node.value { | |
// Node to remove is on left | |
parentNode = node | |
currentNode = node.left | |
} else if value > node.value { | |
// Node to remove is on right | |
parentNode = node | |
currentNode = node.right | |
} else { | |
// Found node to remove | |
if let left = node.left, let right = node.right { | |
// Node has two children | |
node.value = right.getMinValue() | |
right.remove(value: node.value, parentNode: node) | |
} else if parentNode == nil { | |
// Node to remove is root | |
if let left = node.left { | |
node.value = left.value | |
node.left = left.left | |
node.right = left.right | |
} else if let right = node.right { | |
node.value = right.value | |
node.left = right.left | |
node.right = right.right | |
} else { | |
// Single node tree; do nothing | |
} | |
} else if let parent = parentNode { | |
// Node to remove has a parent | |
if let parentLeft = parent.left, parentLeft === node { | |
if let left = node.left { | |
parent.left = left | |
} else { | |
parent.left = node.right | |
} | |
} else if let parentRight = parent.right, parentRight === node { | |
if let left = node.left { | |
parent.right = left | |
} else { | |
parent.right = node.right | |
} | |
} | |
} | |
break | |
} | |
} | |
/* | |
Cases: | |
- Node has no children. | |
- Node has left child. | |
- Node has right child. | |
- Node has left and right child. | |
Special Cases: | |
- Duplicate values | |
- value == nil && parentNode == nil, return root | |
- value != self.value && parentNode == nil, invalid, return root | |
- value == self.value && parentNode == nil, removal of root | |
- value == parentNode.left.value || value == parentNode.right.value | |
Utility Functions: | |
- getMinValue(), return min value in subtree | |
Logic: | |
- Problematic logic with two children. | |
- Logic dependent on other logic causing confusing logic. | |
*/ | |
return self | |
} | |
// Returns minimum value in branch | |
func getMinValue() -> Int { | |
var currentNode: BST? = self | |
while currentNode?.left != nil { | |
currentNode = currentNode!.left | |
} | |
return currentNode!.value | |
} | |
} | |
} | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment