Trie-generic-lowercase

import Foundation

protocol Alphabet {
    associatedtype AlphabetCharacter: Comparable, Hashable

    func allCharacters() -> [AlphabetCharacter]

    init()

}


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
        }
    }

}



let myTrie = Trie<LowercaseLetters>()

myTrie.insert("to");
myTrie.insert("tour");
myTrie.insert("tea");
myTrie.insert("roll");
myTrie.insert("round");

myTrie.remove("roLl")
print (myTrie.query("roll"))      // true

myTrie.removeWordsPrefix("r")
print (myTrie.query("roll"))      // false
print(myTrie.wordsSamePrefix("t"))  // 3

// current words in our dictionary :tea, to, tour
let word: String = myTrie.kthWord(atIndex: 2)!
print(word) // tour