Instantly share code, notes, and snippets.

Embed
What would you like to do?
import Foundation
protocol Alphabet {
associatedtype AlphabetCharacter: Comparable, Hashable
func allCharacters() -> [AlphabetCharacter]
init()
}
struct Stack<Element> {
var items = [Element]()
mutating func push(_ item: Element) {
items.append(item)
}
mutating func pop() -> Element {
return items.removeLast()
}
func isEmpty() -> Bool {
return items.count == 0
}
}
protocol CharacterSequence {
associatedtype SequenceCharacter: Comparable, Hashable
func toSequence() -> [SequenceCharacter]
init(characters: [SequenceCharacter])
}
extension String: CharacterSequence {
typealias SequenceCharacter = Character
func toSequence() -> [SequenceCharacter] {
return Array(characters)
}
init(characters: [SequenceCharacter]) {
self.init(characters)
}
}
extension Int: CharacterSequence {
typealias SequenceCharacter = Int8
func toSequence() -> [SequenceCharacter] {
return String(self).characters.map { character in
let asString = "\(character)"
let asInt = Int(asString)!
let asInt8 = Int8(asInt)
return asInt8
}
}
init(characters: [SequenceCharacter]) {
var value = 0
var p = 1
for char in characters.reversed() {
value += p * Int(char)
p *= 10
}
self.init(value)
}
}
// This defines An alphabet with the latin leters from a to z
class LowercaseLetters: Alphabet {
typealias AlphabetCharacter = Character
func allCharacters() -> [AlphabetCharacter] {
return Array("abcdefghijklmnopqrstuvwxyz".characters)
}
required init() {
}
}
class DigitsAlphabet: Alphabet {
typealias AlphabetCharacter = Int8
func allCharacters() -> [AlphabetCharacter] {
return Array(0...9)
}
required init() {
}
}
class TrieNode<T> where T:Alphabet {
typealias character = T.AlphabetCharacter
var children: [character:TrieNode<T>] = [:]
var prefixCount: Int = 0
var isFinal:Bool = false
init() {
}
///// called from a PARENT
func createChildFor(_ character: character) -> TrieNode<T> {
let node = TrieNode<T>()
children[character] = node
return node
}
func getOrCreateChildFor(_ character: character) -> TrieNode<T> {
if let child = children[character] {
return child
} else {
return createChildFor(character)
}
}
}
class Trie<T> where T:Alphabet {
typealias character = T.AlphabetCharacter
typealias Node = TrieNode<T>
var root = Node()
func validate(characters: [character]) -> Bool {
let allCharacters = T().allCharacters()
return characters.reduce(true) { result, char in
return result && allCharacters.contains(char)
}
}
func insert<U:CharacterSequence>(_ word: U) where U.SequenceCharacter == character {
insert(characters: word.toSequence())
}
func insert(characters: [character]) {
guard validate(characters: characters) else {
return
}
var node = root
node.prefixCount += 1
for character in characters {
node = node.getOrCreateChildFor(character)
node.prefixCount += 1
}
node.isFinal = true
}
func getNodeFor(characters: [character]) -> Node? {
var node: Node? = root
for character in characters {
node = node?.children[character]
}
return node
}
func query<U:CharacterSequence>(_ word: U) -> Bool where U.SequenceCharacter == character {
return query(characters: word.toSequence())
}
func query(characters: [character]) -> Bool {
guard validate(characters: characters) else {
return false
}
let node = getNodeFor(characters: characters)
if node == nil {
return false
}
return node!.isFinal
}
func getNodeSequence(characters: [character]) -> [(char: character, node: Node)] {
var node: Node? = root
var nodes:[(character,Node)] = []
let tup = (characters[0],node!)
nodes.append(tup)
for character in characters {
node = node?.children[character]
if node == nil {
return nodes
}
let tup = (character, node!)
nodes.append(tup)
}
return nodes
}
func remove<U:CharacterSequence>(_ word: U) where U.SequenceCharacter == character {
remove(characters: word.toSequence())
}
func remove(characters: [character]) {
guard validate(characters: characters) else {
return
}
let charactersNodes = getNodeSequence(characters: characters)
if charactersNodes.count != characters.count + 1 {
return
}
var parent = root
for (char, node) in charactersNodes {
node.prefixCount -= 1
if node.prefixCount == 0 && node !== root {
parent.children.removeValue(forKey: char)
break // we deleted the reference to an entire sequence of nodes
// swift will delete those nodes from memory
}
parent = node
}
charactersNodes.last?.node.isFinal = false
}
func removeWordsPrefix<U:CharacterSequence>(_ word: U) where U.SequenceCharacter == character {
removeWordsPrefix(characters: word.toSequence())
}
func removeWordsPrefix(characters: [character]) {
guard validate(characters: characters) else {
return
}
let charactersNodes = getNodeSequence(characters: characters)
if charactersNodes.count != characters.count + 1 {
return
}
let decreasePrefix = charactersNodes.last!.node.prefixCount
var parent = root
for (char, node) in charactersNodes {
node.prefixCount -= decreasePrefix
if node.prefixCount == 0 && node !== root {
parent.children.removeValue(forKey: char)
break // we deleted the reference to an entire sequence of nodes
// swift will delete those nodes from memory
}
parent = node
}
}
func kthWord<U:CharacterSequence>(atIndex index: Int, from node: Node? = nil, characters: [character] = []) -> U? where U.SequenceCharacter == character {
let node = node ?? root
if index == 0 && node.isFinal {
return U(characters: characters)
}
var skipped = 0
if node.isFinal {
skipped += 1
}
for char in T().allCharacters() {
if let child = node.children[char] {
if skipped + child.prefixCount > index {
// find the word in the current child
return kthWord(atIndex: index - skipped, from: child, characters: characters + [char])
} else {
// skip words from the current child
skipped += child.prefixCount
}
}
}
return nil
}
func wordsSamePrefix<U:CharacterSequence>(_ word: U) -> Int where U.SequenceCharacter == character {
return wordsSamePrefix(characters: word.toSequence())
}
func wordsSamePrefix(characters: [character]) -> Int {
guard validate(characters: characters) else {
return 0
}
if let node = getNodeFor(characters: characters) {
return node.prefixCount
} else {
return 0
}
}
func suggestWords<U:CharacterSequence>(_ word: U) -> [U] where U.SequenceCharacter == character {
return suggestWords(characters: word.toSequence())
}
func suggestWords<U:CharacterSequence>(characters: [character]) -> [U] where U.SequenceCharacter == character{
guard validate(characters: characters) else {
return []
}
let lastNodeCharacters = getNodeFor(characters: characters)
if lastNodeCharacters == nil {
return []
}
var discoveredNodes:[(node: Node, char: character, counter: Int)] = DFS(node: lastNodeCharacters!, char: characters[characters.count - 1])
// discoveredNodes also includes the last character from characters parameter
// it can easily be removed here, since we know it's the first discovered node
discoveredNodes.remove(at: 0)
return buildSeq(characters: characters, discoveredNodes: &discoveredNodes)
}
// parameters: node and corresponding character
// counter == prefixCount; useful when building the sequence
func DFS(node: Node, char: character ) -> [(node: Node, char: character, counter: Int)] {
let allCharacters = T().allCharacters()
var discoveredNodes:[(node: Node, char: character, counter: Int)] = []
var stackOfNodes = Stack<(node: Node,char: character, counter: Int)>()
stackOfNodes.push((node: node, char: char, counter: node.prefixCount ))
var nodeChar:(node: Node, char: character, counter: Int)
while !stackOfNodes.isEmpty(){
nodeChar = stackOfNodes.pop()
discoveredNodes.append(nodeChar)
for character in allCharacters {
if let inter = nodeChar.node.children[character] {
stackOfNodes.push((node:inter, char:character, counter:inter.prefixCount))
}
}
}
return discoveredNodes
}
// takes the prefix, the nodes from DFS and builds words (suggestions)
func buildSeq<U:CharacterSequence>(characters: [character], discoveredNodes: inout [(node: Node, char: character, counter: Int)]) -> [U] where U.SequenceCharacter == character {
var res:[U] = []
for i in 0..<discoveredNodes.count {
if discoveredNodes[i].node.isFinal == true {
var w:[character] = []
// counter tells us which nodes characters are to be included in the sequence
// if counter > 0 then the node's character belongs to the sequence and we decrement the counter
// otherwise, the node was used in other sequences and was "used up"
print("building seq starting with:",characters)
for j in 0...i{
if discoveredNodes[j].counter > 0 {
print(discoveredNodes[j].char)
w.append(discoveredNodes[j].char)
discoveredNodes[j].counter -= 1
}
}
w = characters + w
res.append(U(characters: w))
}
}
return res
}
}
let myTrie = Trie<LowercaseLetters>()
myTrie.insert("to")
myTrie.insert("tour")
myTrie.insert("tea")
myTrie.insert("roll")
myTrie.insert("round")
myTrie.insert("route")
print (myTrie.suggestWords("ro"))
// let myTrie2 = Trie<DigitsAlphabet>()
// myTrie2.insert(12);
// myTrie2.insert(123);
// myTrie2.insert(13);
// myTrie2.insert(133);
// myTrie2.insert(134);
// myTrie2.insert(211);
//
// myTrie2.remove(12)
// print(myTrie2.query(12)) // false
// print(myTrie2.query(123)) // true
//
// myTrie2.removeWordsPrefix(12)
// print(myTrie2.query(123)) // false
// print(myTrie2.wordsSamePrefix(13)) // 3 : 13, 133, 134
//
// print (myTrie2.suggestWords(13))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment