Useful data structures in Swift 2.2

DataStructures.swift

/**
 Hi
 */

HashMap.swift

/**
 A basic dictionary type.
 
 Features:
 - O(1) insertion, lookup, and removal
 */
class HashMap<K: Hashable, V> {

    typealias Pair = (key: K, value: V)
    typealias Bucket = LinkedList<Pair>

    private var buckets: [Bucket]

    private(set) var count: Int = 0

    private let maximumLoad: Float = 1.0
    var load: Float {
        return Float(count) / Float(buckets.count)
    }

    var pairs: [Pair] {
        // Flatten array of buckets (array of linked lists) into an array of pairs
        return buckets.map({ bucket in
            bucket.map { $0 }
        }).reduce([]) { $0 + $1 }
    }

    private class func makeBuckets(count: Int) -> [Bucket] {
        var buckets: [Bucket] = []
        for _ in 0..<count {
            buckets.append(Bucket())
        }
        return buckets
    }

    init(bucketCount: Int = 8) {
        buckets = HashMap.makeBuckets(bucketCount)
    }

    private func bucketIndex(key: K) -> Int {
        return key.hashValue & (buckets.count - 1)
    }

    private func rehash() {
        let tmp = pairs

        buckets = HashMap.makeBuckets(buckets.count * 2)
        count = 0

        for p in tmp {
            self[p.key] = p.value
        }
    }

    func updateValue(value: V, forKey key: K) -> V? {
        if load >= maximumLoad { rehash() }

        var index: Bucket.Index?
        var old: Pair?

        let bucket = buckets[bucketIndex(key)]
        for (i, p) in bucket.enumerate() {
            if p.key == key {
                index = i
                old = p
                break
            }
        }

        if let i = index {
            bucket.remove(atIndex: i) // O(i/2)
            count -= 1
        }

        bucket.append((key: key, value: value))
        count += 1

        return old?.value
    }

    func removeValueForKey(key: K) -> V? {
        var index: Bucket.Index?
        var old: Pair?

        let bucket = buckets[bucketIndex(key)]
        for (i, p) in bucket.enumerate() {
            if p.key == key {
                index = i
                old = p
                break
            }
        }

        if let i = index, o = old {
            bucket.remove(atIndex: i) // O(i/2)
            count -= 1
            return o.value
        } else {
            return nil
        }
    }

    subscript(key: K) -> V? {
        get {
            let bucket = buckets[bucketIndex(key)]
            for pair in bucket {
                if pair.key == key {
                    return pair.value
                }
            }
            return nil
        }
        set(newValue) {
            if let v = newValue {
                updateValue(v, forKey: key)
            } else {
                removeValueForKey(key)
            }
        }
    }
    
}

extension HashMap: CustomStringConvertible {

    var description: String {
        let s: String = pairs.reduce("\n") { (soFar, pair) in
            "\(soFar)\(pair.key): \(pair.value)\n"
        }
        return "{\(s)}"
    }

}

extension HashMap: CustomDebugStringConvertible {

    var debugDescription: String {
        return "\(description)\ncount: \(count), buckets: \(buckets.count), load: \(load)"
    }

}

extension HashMap: SequenceType {

    typealias Generator = AnyGenerator<Pair>

    func generate() -> Generator {
        return Generator(pairs.generate())
    }

}

LinkedList.swift

/**
 Swift doesn't allow generic types to contain nested types, which means
 that even though this really belongs inside of the LinkedList class
 definition, it has to be top-level.
 */
class LinkedListNode<T>: CustomStringConvertible {

    let value: T

    var prev: LinkedListNode?
    var next: LinkedListNode?

    init(_ value: T) {
        self.value = value
    }

    var description: String {
        var s: String!
        if let csc = value as? CustomStringConvertible {
            s = csc.description
        } else {
            s = String(value)
        }
        return "( \(s) )"
    }

}

/**
 A doubly-linked list of values of type T.
 
 Features:
 - O(1) access to and removal of head/tail
 - O(1) prepend and append
 - O(n/2) internal access and insertion/removal
 - O(n) reversal
 */
class LinkedList<T> {

    typealias Node = LinkedListNode<T>

    private(set) var head: Node?
    private(set) var tail: Node?
    private(set) var count: Int = 0

    init() {}

    private func added(n: Node) {
        count += 1
        if count == 1 {
            head = n
            tail = n
        }
    }

    private func removed(n: Node) {
        count -= 1
        if n === head {
            head = n.next
        } else if n === tail {
            tail = n.prev
        }
    }

    /**
     Traverse the list as efficiently as possible to find the node at the given
     index. That just means starting at the head for indexes < count / 2 and
     starting at the tail for indexes >= count / 2.
     
     O(1) for index ∈ {0, count}
     O(n/2) otherwise
     */
    func nodeAt(index: Int) -> Node? {
        guard index >= 0 && index < count else { return nil }

        var n: Node?
        if index < count / 2 {
            n = head
            for _ in 0..<index {
                n = n!.next
            }
        } else {
            n = tail
            for _ in (count - 1).stride(to: index, by: -1) {
                n = n!.prev
            }
        }

        return n
    }

    /**
     O(1)
     */
    func prepend(value: T) {
        let n = Node(value)

        n.next = head
        if let h = head {
            h.prev = n
        }
        head = n

        added(n)
    }

    /**
     O(1)
     */
    func append(value: T) {
        let n = Node(value)

        n.prev = tail
        if let t = tail {
            t.next = n
        }
        tail = n
        
        added(n)
    }

    /**
     O(1) for index ∈ {0, count}
     O(n/2) otherwise
     */
    func insert(value: T, atIndex index: Int) {
        guard index >= 0 && index <= count else { return }

        let n = Node(value)

        if index == 0 {
            prepend(value)
        } else if index == count {
            append(value)
        } else {
            guard let tmp = nodeAt(index) else { return }
            if let prev = tmp.prev {
                prev.next = n
                n.prev = prev
            }
            tmp.prev = n
            n.next = tmp
        }

        added(n)
    }

    /**
     O(1) for index ∈ {0, count}
     O(n/2) otherwise
     */
    func remove(atIndex index: Int) -> T? {
        guard let n = nodeAt(index) else { return nil }

        if let prev = n.prev {
            prev.next = n.next
        }
        if let next = n.next {
            next.prev = n.prev
        }

        removed(n)
        return n.value
    }

    /**
     O(n)
     */
    func reverse() {
        var tmp: Node?
        var node = head

        while let n = node {
            tmp = n.prev
            n.prev = n.next
            n.next = tmp
            node = n.prev
        }

        tmp = head
        head = tail
        tail = tmp
    }

}

extension LinkedList: CustomStringConvertible {

    var description: String {
        var nodes: [Node] = []

        var curr: LinkedListNode! = head
        while curr != nil {
            nodes.append(curr)
            curr = curr.next
        }

        let valuesString = nodes.map({ $0.description }).joinWithSeparator(" <=> ")
        return "{ \(valuesString) }"
    }

}

extension LinkedList: CustomDebugStringConvertible {

    var debugDescription: String {
        var s = "\ncount: \(count)"
        if let h = head {
            s += ", head: \(h.description)"
        }
        if let t = tail {
            s += ", tail: \(t.description)"
        }
        return "\(description)\n\(s)"
    }

}

extension LinkedList: SequenceType {

    func generate() -> AnyGenerator<T> {
        var curr: LinkedListNode? = head
        return AnyGenerator {
            let value = curr?.value
            curr = curr?.next
            return value
        }
    }

}

extension LinkedList: CollectionType {

    typealias Index = Int

    var startIndex: Index { return 0 }
    var endIndex: Index { return count }

    subscript(i: Index) -> T {
        return nodeAt(i)!.value
    }

}