Create a gist now

Instantly share code, notes, and snippets.

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