Skip to content

Instantly share code, notes, and snippets.

Created June 7, 2014 07:23
Show Gist options
  • Save anonymous/b76f3500c4d7a2a35401 to your computer and use it in GitHub Desktop.
Save anonymous/b76f3500c4d7a2a35401 to your computer and use it in GitHub Desktop.
//
// Trie.swift
// SceneTest
//
// Created by Greg Titus on 6/7/14.
// Copyright (c) 2014 The Omni Group. All rights reserved.
//
// This is essentially the same data structure as OFTrie in OmniFoundation, except that the OmniFoundation version is mutable and ephemeral,
// and this version is immutable and persistent. Every time you add a string to the trie, it replaces all the nodes along the string's path
// to the root, and the only mutable change is reassigning the root. Old copies of the trie are unaffected.
// 100 or so lines of code here, 600 or so in 3 classes (6 .h or .m files) in the Objective-C version.
// some conveniences, to treat the Sliceable UnicodeScalarView of a String as if it's a sort of cons'd list type...
extension String.UnicodeScalarView {
var isEmpty : Bool { return self.startIndex == self.endIndex }
var head : UnicodeScalar { return self[self.startIndex] }
var tail : String.UnicodeScalarView { return self[self.startIndex.succ() .. self.endIndex] }
}
struct TrieNode {
enum TrieChildren {
// either a single leaf value, with a trailing suffix
case Leaf(String.UnicodeScalarView)
// or a branch, mapping first character to child node
case Branch(Dictionary<UnicodeScalar, TrieNode>)
}
var value : AnyObject?
var children : TrieChildren
func insertValue(newValue: AnyObject, forChars: String.UnicodeScalarView) -> TrieNode {
switch children {
case .Leaf(let suffix):
if (suffix == forChars) {
return TrieNode(value: newValue, children: .Leaf(forChars))
} else if (forChars.isEmpty) {
let oldHead = suffix.head
let oldLeaf = TrieNode(value: value, children: .Leaf(suffix.tail))
return TrieNode(value: newValue, children: .Branch([oldHead: oldLeaf]))
} else if (suffix.isEmpty) {
let newHead = forChars.head
let newLeaf = TrieNode(value: newValue, children: .Leaf(forChars.tail))
return TrieNode(value: value, children: .Branch([newHead: newLeaf]))
} else {
let oldHead = suffix.head
let newHead = forChars.head
let oldLeaf = TrieNode(value: value, children: .Leaf(suffix.tail))
if (oldHead == newHead) {
let newChild = oldLeaf.insertValue(newValue, forChars: forChars.tail)
return TrieNode(value: nil, children: .Branch([newHead : newChild]))
} else {
let newLeaf = TrieNode(value: newValue, children: .Leaf(forChars.tail))
return TrieNode(value: nil, children: .Branch([oldHead: oldLeaf, newHead: newLeaf]))
}
}
case .Branch(let subnodes):
if (forChars.isEmpty) {
return TrieNode(value: newValue, children: children)
}
var newChild : TrieNode
let charHead = forChars.head
if let child = subnodes[charHead] {
newChild = child.insertValue(newValue, forChars: forChars.tail)
} else {
newChild = TrieNode(value: newValue, children: .Leaf(forChars.tail))
}
var newSubnodes = subnodes
newSubnodes[charHead] = newChild
return TrieNode(value: value, children: .Branch(newSubnodes))
}
}
func lookup (str : String.UnicodeScalarView) -> AnyObject? {
if (str.isEmpty) {
return value
}
switch children {
case .Leaf(let suffix):
return suffix == str ? value : nil
case .Branch(let subnodes):
return subnodes[str.head]?.lookup(str.tail)
}
}
}
class Trie {
var root : TrieNode? = nil
func insert(value: AnyObject, forString: String) {
if root {
root = root!.insertValue(value, forChars: forString.unicodeScalars)
} else {
root = TrieNode(value: value, children: .Leaf(forString.unicodeScalars))
}
}
func lookup(str : String) -> AnyObject? {
return root?.lookup(str.unicodeScalars)
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment