Trie in Swift

gistfile1.swift

// Playground - noun: a place where people can play

import Foundation

extension String {
	subscript (i: Int) -> String {
		return String(Array(self)[i])
	}
	
	func count() -> Int {
		return countElements(self)
	}
	
	subscript (r: Range<Int>) -> String {
		var start = advance(startIndex, r.startIndex)
		var end = advance(startIndex, r.endIndex)
		return substringWithRange(Range(start: start, end: end))
	}
}

extension Array {
	func each(doThis: (element: T) -> Void) {
		for e in self {
			doThis(element: e)
		}
	}
}

extension Int {
	
	func times(closure:(Int)->Void) {
		for i in 0..<self {
			closure(i)
		}
	}
	
}

class TrieNode {
	var letter:String?
	var children:[TrieNode] = []
	var level:Int = 0
	var isWord:Bool = false
	var fullWord:String?
	
	init(){}
}

func setupTrie(var input:String, inout rootNode:TrieNode?, var currentIndex:Int = 0, var level:Int = 0) {
	
	if (rootNode == nil) {
		rootNode = TrieNode()
		rootNode?.letter = input[0]
	}
	
	++level
	var foundChild = false
	if let children = rootNode?.children {
		children.count.times({ i in
			var child:TrieNode? = children[i]
			if (input[currentIndex] == child!.letter) {
				foundChild = true
				if (input.count() > currentIndex) {
					setupTrie(input, &child, currentIndex: ++currentIndex)
				}
			}
		})
	}
	
	if (!foundChild) {
		
		var child:TrieNode? = TrieNode()
		child?.letter = input[currentIndex]
		child?.level = level
		child?.fullWord = input[0...currentIndex]
		if (input.count() - 1 == currentIndex) {
			child?.isWord = true
		}
		
		rootNode?.children.append(child!)
		++currentIndex
		
		if (input.count() > currentIndex) {
			setupTrie(input, &child, level: level, currentIndex: currentIndex)
		}
	}

}

var rootNode:TrieNode? = TrieNode()
var words:[String] = ["Base", "Bass", "Baseball", "Basely"]

words.each { (word:String) -> Void in
	setupTrie(word, &rootNode)
}

func searchForWords(rootNode:TrieNode?, var queue:[TrieNode] = [], var foundWords:[TrieNode] = []) -> [TrieNode] {
	
	if (rootNode?.children.count == 0) {
		return foundWords
	}

	rootNode?.children.each({ (element:TrieNode) in
		if (element.isWord) {
			foundWords.append(element)
		}
		queue.append(element)
		
		var item = queue.removeAtIndex(queue.count - 1)
		foundWords = searchForWords(item, queue: queue, foundWords: foundWords)
	})

	return foundWords
}

func searchTrie(rootNode:TrieNode?, var targetWord:String, var currentIndex:Int = 0, var foundNode:TrieNode? = nil) -> TrieNode? {
	
	if let children = rootNode?.children {
		if (currentIndex == targetWord.count()) {
			return rootNode!
		} else {
			
			children.count.times({ i in
				var child:TrieNode? = children[i]
				if (currentIndex >= targetWord.count() || targetWord[currentIndex] == child!.letter) {
					var node:TrieNode? = searchTrie(child, targetWord, currentIndex: ++currentIndex)
					if (node != nil) {
						foundNode = node
					}
				}
			})
		}
	}
	
	return foundNode
}

var node:TrieNode? = searchTrie(rootNode, "Base")
searchForWords(node).sorted { (a:TrieNode, b:TrieNode) -> Bool in
	a.level < b.level
}.each { (element) -> Void in
	println(element.fullWord)
}