Red Black Tree with indirect and pattern matching for Swift 2.0b4
enum Color { case R, B } | |
indirect enum Tree<Element: Comparable> { | |
case Empty | |
case Node(Color,Tree<Element>,Element,Tree<Element>) | |
init() { self = .Empty } | |
init(_ x: Element, color: Color = .B, | |
left: Tree<Element> = .Empty, right: Tree<Element> = .Empty) | |
{ | |
self = .Node(color, left, x, right) | |
} | |
} | |
extension Tree { | |
func contains(x: Element) -> Bool { | |
switch self { | |
case .Empty: return false | |
case let .Node(_, left, y, right): | |
if x < y { return left.contains(x) } | |
if y < x { return right.contains(x) } | |
return true | |
} | |
} | |
} | |
private func balance<T>(tree: Tree<T>) -> Tree<T> { | |
switch tree { | |
case let .Node(.B, .Node(.R, .Node(.R, a, x, b), y, c), z, d): | |
return .Node(.R, .Node(.B,a,x,b),y,.Node(.B,c,z,d)) | |
case let .Node(.B, .Node(.R, a, x, .Node(.R, b, y, c)), z, d): | |
return .Node(.R, .Node(.B,a,x,b),y,.Node(.B,c,z,d)) | |
case let .Node(.B, a, x, .Node(.R, .Node(.R, b, y, c), z, d)): | |
return .Node(.R, .Node(.B,a,x,b),y,.Node(.B,c,z,d)) | |
case let .Node(.B,a,x,.Node(.R,b,y,.Node(.R,c,z,d))): | |
return .Node(.R, .Node(.B,a,x,b),y,.Node(.B,c,z,d)) | |
default: | |
return tree | |
} | |
} | |
private func ins<T>(into: Tree<T>, _ x: T) -> Tree<T> { | |
switch into { | |
case .Empty: return Tree(x, color: .R) | |
case let .Node(c, l, y, r): | |
if x < y { return balance(Tree(y, color: c, left: ins(l,x), right: r)) } | |
if y < x { return balance(Tree(y, color: c, left: l, right: ins(r, x))) } | |
return into | |
} | |
} | |
extension Tree { | |
func insert(x: Element) -> Tree { | |
switch ins(self, x) { | |
case let .Node(_,l,y,r): | |
return .Node(.B,l,y,r) | |
default: | |
fatalError("ins should never return an empty tree") | |
} | |
} | |
} | |
extension Tree: SequenceType { | |
func generate() -> AnyGenerator<Element> { | |
// stack-based iterative inorder traversal to | |
// make it easy to use with anyGenerator | |
var stack: [Tree] = [] | |
var current: Tree = self | |
return anyGenerator { _ -> Element? in | |
while true { | |
switch current { | |
case let .Node(_,l,_,_): | |
stack.append(current) | |
current = l | |
case .Empty where !stack.isEmpty: | |
switch stack.removeLast() { | |
case let .Node(_,_,x,r): | |
current = r | |
return x | |
default: | |
break | |
} | |
case .Empty: | |
return nil | |
} | |
} | |
} | |
} | |
} | |
extension Tree: ArrayLiteralConvertible { | |
init<S: SequenceType where S.Generator.Element == Element>(_ source: S) { | |
self = source.reduce(Tree()) { $0.insert($1) } | |
} | |
init(arrayLiteral elements: Element...) { | |
self = Tree(elements) | |
} | |
} | |
import Darwin | |
extension Array { | |
func shuffle() -> [Element] { | |
var list = self | |
for i in 0..<(list.count - 1) { | |
let j = Int(arc4random_uniform(UInt32(list.count - i))) + i | |
swap(&list[i], &list[j]) | |
} | |
return list | |
} | |
} | |
let engines = [ | |
"Daisy", "Salty", "Harold", "Cranky", | |
"Thomas", "Henry", "James", "Toby", | |
"Belle", "Diesel", "Stepney", "Gordon", | |
"Captain", "Percy", "Arry", "Bert", | |
"Spencer", | |
] | |
// test various inserting engines in various different permutations | |
for permutation in [engines, engines.sort(), engines.sort(>), engines.shuffle(),engines.shuffle(),engines.shuffle()] { | |
let t = Tree(permutation) | |
assert(!t.contains("Fred")) | |
assert(t.elementsEqual(t.insert("Thomas"))) | |
assert(!engines.contains { !t.contains($0) }) | |
assert(t.elementsEqual(engines.sort())) | |
print(",".join(t)) | |
} | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment