Trie - lowercase

import Foundation

class TrieNode{

    var children: [Character:TrieNode] = [:]
    var isFinal: Bool = false
    var prefixCount: Int = 0


    init() {

    }

///// called from a PARENT
    func createChildFor(_ character: Character) -> TrieNode {
        let node = TrieNode()

        children[character] = node

        return node
    }

    func getOrCreateChildFor(_ character: Character) -> TrieNode {
        if let child = children[character] {
            return child
        } else {
            return createChildFor(character)
        }
    }
}
///////////

class Trie {
    var root = TrieNode()

    func validate(characters: [Character]) -> Bool {
        var alphabetString = "abcdefghijklmnopqrstuvxyz"
        let alphabetSequence = Array(alphabetString.characters)

        return characters.reduce(true) { result, char in
            return result && alphabetSequence.contains(char)
        }
    }

    func getNodeFor(_ characters: [Character]) -> TrieNode? {
        var node: TrieNode? = root
        for character in characters {
            node = node?.children[character]
        }
        return node
    }

    func insert(_ word: String){
        insert(characters: Array(word.characters))
    }

    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 query(_ word: String) -> Bool {
        return query(characters: Array(word.characters))
    }

    func query(characters: [Character]) -> Bool {
      guard validate(characters: characters) else {
          return false
      }

      let node = getNodeFor(characters)

      if node == nil {
        return false
      }
      return node!.isFinal
    }

    func remove(_ word: String){
        remove(characters: Array(word.characters))
    }

    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 getNodeSequence(characters: [Character]) -> [(char: Character, node: TrieNode)] {
      var node : TrieNode? = root
      var nodes:[(Character,TrieNode)] = []

      let tup = (Character("@"),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 wordsSamePrefix(_ word: String) -> Int {
        return wordsSamePrefix(characters: Array(word.characters))
    }

    func wordsSamePrefix(characters: [Character]) -> Int {
      guard validate(characters: characters) else {
          return 0
      }

      let node = getNodeFor(characters)

      if node == nil {
        return 0
      }

      return node!.prefixCount
    }


    func removeWordsPrefix(_ word: String)  {
       removeWordsPrefix(characters: Array(word.characters))
    }

    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(atIndex index: Int, from node: TrieNode? = nil, characters: [Character] = []) -> String? {
        let node = node ?? root

        if index >= self.uniqueWords() {
          return nil
        }

        var alphabetString = "abcdefghijklmnopqrstuvxyz"
        let alphabetSequence = Array(alphabetString.characters)

        if index == 0 && node.isFinal {
            return String(characters)
        }

        var skipped = 0
        if node.isFinal {
          skipped += 1
        }

        for char in alphabetSequence {
            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 uniqueWords() -> Int {
      return root.prefixCount
    }
  }



  let myTrie = Trie()

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

  let word = myTrie.kthWord(atIndex:2)
  print(word!)   // optional "tea"

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

  myTrie.removeWordsPrefix("r")
  print (myTrie.query("round"))  //false

  print(myTrie.wordsSamePrefix("t")) //3