Skip to content

Instantly share code, notes, and snippets.

@adamrothman
Last active March 25, 2016 00:37
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save adamrothman/5ee0f55307432135284a to your computer and use it in GitHub Desktop.
Save adamrothman/5ee0f55307432135284a to your computer and use it in GitHub Desktop.
Useful data structures in Swift 2.2
/**
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())
}
}
/**
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
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment