Created
May 11, 2015 15:50
-
-
Save airspeedswift/8f8b8548737bb1b8a0d0 to your computer and use it in GitHub Desktop.
Swift Ordered Set based on Red/Black tree
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
enum Color { | |
case Red, Black | |
} | |
class Tree<T: Comparable> { | |
let value: T | |
let left: Tree<T>? | |
let right: Tree<T>? | |
let color: Color | |
init(_ color: Color, _ value: T, left: Tree<T>? = nil, right: Tree<T>? = nil) { | |
self.value = value | |
self.left = left | |
self.right = right | |
self.color = color | |
} | |
} | |
func balance<T>(tree: Tree<T>) -> Tree<T> { | |
if tree.color == .Black { | |
if let x = tree.left, y = x.right where y.color == .Red && x.color == .Red { | |
let z = tree, a = x.left, b = y.left, c = y.right, d = z.right | |
return Tree(.Red, y.value, | |
left: Tree(.Black, x.value, left: a, right: b), | |
right: Tree(.Black, z.value, left: c, right: d)) | |
} | |
if let y = tree.left, x = y.left where y.color == .Red && x.color == .Red { | |
let z = tree, a = x.left, b = x.right, c = y.right, d = z.right | |
return Tree(.Red, y.value, | |
left: Tree(.Black, x.value, left: a, right: b), | |
right: Tree(.Black, z.value, left: c, right: d)) | |
} | |
if let y = tree.right, z = y.right where y.color == .Red && z.color == .Red { | |
let x = tree, a = x.left, b = y.left, c = z.left, d = z.right | |
return Tree(.Red, y.value, | |
left: Tree(.Black, x.value, left: a, right: b), | |
right: Tree(.Black, z.value, left: c, right: d)) | |
} | |
if let z = tree.right, y = z.left where y.color == .Red && z.color == .Red { | |
let x = tree, a = x.left, b = y.left, c = y.right, d = z.right | |
return Tree(.Red, y.value, | |
left: Tree(.Black, x.value, left: a, right: b), | |
right: Tree(.Black, z.value, left: c, right: d)) | |
} | |
} | |
return tree | |
} | |
func insert<T>(into: Tree<T>?, value: T) -> Tree<T> { | |
let t = ins(into, value) | |
return Tree(.Black, t.value, left: t.left, right: t.right) | |
} | |
func ins<T>(into: Tree<T>?, value: T) -> Tree<T> { | |
switch into { | |
case .None: | |
return Tree(.Red, value) | |
case .Some(let t) where value < t.value: | |
return balance(Tree(t.color, t.value, left: ins(t.left, value), right: t.right)) | |
case .Some(let t) where value > t.value: | |
return balance(Tree(t.color, t.value, left: t.left, right: ins(t.right, value))) | |
case .Some(let t): | |
return t | |
} | |
} | |
func contains<T>(tree: Tree<T>?, member: T) -> Bool { | |
switch tree { | |
case .None: | |
return false | |
case .Some(let t) where member < t.value: | |
return contains(t.left, member) | |
case .Some(let t) where member > t.value: | |
return contains(t.right, member) | |
default: | |
return true | |
} | |
} | |
func inorderSequence<T>(tree: Tree<T>?) -> SequenceOf<T> { | |
return SequenceOf { _ -> GeneratorOf<T> in | |
var stack: [Tree<T>] = [] | |
var current = tree | |
return GeneratorOf { | |
while true { | |
if let t = current { | |
stack.append(t) | |
current = t.left | |
} | |
else if !stack.isEmpty { | |
let pop = stack.removeLast() | |
current = pop.right | |
return pop.value | |
} | |
else { | |
return nil | |
} | |
} | |
} | |
} | |
} | |
let engines = [ | |
"Thomas", "Henry", "James", "Toby", | |
"Belle", "Diesel", "Stepney", "Gordon", | |
"Daisy", "Salty", "Harold", "Cranky", | |
"Captain", "Percy", "Arry", "Bert", | |
] | |
let t = reduce(engines, nil, insert) | |
", ".join(inorderSequence(t)) | |
for permutation in [engines, sorted(engines,<), sorted(engines,>), dropFirst(engines) + [engines[0]]] { | |
let t = reduce(permutation, nil, insert) | |
assert(!contains(t, "Fred")) | |
assert(reduce(permutation, true, { $0.0 && contains(t, $0.1) }).0) | |
assert(equal(inorderSequence(t), sorted(engines))) | |
} | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment