Skip to content

Instantly share code, notes, and snippets.

@airspeedswift
Created October 4, 2015 00:37
Show Gist options
  • Save airspeedswift/8fd42496679c51b34859 to your computer and use it in GitHub Desktop.
Save airspeedswift/8fd42496679c51b34859 to your computer and use it in GitHub Desktop.
red black tree updated for 2.1b2
private enum ListNode<Element> {
case End
indirect case Node(Element, next: ListNode<Element>)
func cons(x: Element) -> ListNode<Element> {
return .Node(x, next: self)
}
}
public struct ListIndex<Element> {
private let node: ListNode<Element>
private let tag: Int
}
extension ListIndex: ForwardIndexType {
public func successor() -> ListIndex<Element> {
guard case let .Node(_, next: next) = node
else { fatalError("cannot increment endIndex") }
return ListIndex(node: next, tag: tag.predecessor())
}
}
public func == <T>(lhs: ListIndex<T>, rhs: ListIndex<T>) -> Bool {
return lhs.tag == rhs.tag
}
public struct List<Element> {
public typealias Index = ListIndex<Element>
public var startIndex: Index
public var endIndex: Index { return Index(node: .End, tag: 0) }
public init() {
startIndex = Index(node: .End, tag: 0)
}
}
extension List {
public mutating func push(x: Element) {
startIndex = Index(node: startIndex.node.cons(x),
tag: startIndex.tag.successor())
}
public mutating func pop() -> Element? {
guard case let .Node(x, next: _) = startIndex.node
else { return nil }
++startIndex
return x
}
}
extension List: ArrayLiteralConvertible {
public init(arrayLiteral elements: Element...) {
self = List()
for x in elements.reverse() { push(x) }
}
}
private enum Color { case R, B }
private 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 {
guard case let .Node(_,left,y,right) = self
else { return false }
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> {
guard case let .Node(c, l, y, r) = into
else { return Tree(x, color: .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 {
guard case let .Node(_,l,y,r) = ins(self, x)
else { fatalError("ins should never return an empty tree") }
return .Node(.B,l,y,r)
}
}
extension Tree: SequenceType {
func generate() -> AnyGenerator<Element> {
var stack: List<Tree> = []
var current: Tree = self
return anyGenerator { _ -> Element? in
while true {
// if there's a left-hand node, head down it
if case let .Node(_,l,_,_) = current {
stack.push(current)
current = l
}
// if there isn’t, head back up, going right as
// soon as you can:
else if case let .Node(_,_,x,r)? = stack.pop() {
current = r
return x
}
else {
// otherwise, we’re done
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
guard i != j else { continue }
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()] {
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(t.joinWithSeparator(","))
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment