trie.swift

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