Guest User

Trie - generic - complete

a guest
Feb 13th, 2017
83
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1.  
  2. import Foundation
  3.  
  4. protocol Alphabet {
  5.     associatedtype AlphabetCharacter: Comparable, Hashable
  6.  
  7.     func allCharacters() -> [AlphabetCharacter]
  8.  
  9.     init()
  10.  
  11. }
  12.  
  13. struct Stack<Element> {
  14.     var items = [Element]()
  15.     mutating func push(_ item: Element) {
  16.         items.append(item)
  17.     }
  18.     mutating func pop() -> Element {
  19.         return items.removeLast()
  20.     }
  21.  
  22.     func isEmpty() -> Bool {
  23.         return items.count == 0
  24.     }
  25. }
  26.  
  27.  
  28. protocol CharacterSequence {
  29.     associatedtype SequenceCharacter: Comparable, Hashable
  30.     func toSequence() -> [SequenceCharacter]
  31.     init(characters: [SequenceCharacter])
  32. }
  33.  
  34. extension String: CharacterSequence {
  35.     typealias SequenceCharacter = Character
  36.  
  37.     func toSequence() -> [SequenceCharacter] {
  38.         return Array(characters)
  39.     }
  40.  
  41.     init(characters: [SequenceCharacter]) {
  42.         self.init(characters)
  43.     }
  44. }
  45.  
  46. extension Int: CharacterSequence {
  47.      typealias SequenceCharacter = Int8
  48.  
  49.      func toSequence() -> [SequenceCharacter] {
  50.          return String(self).characters.map { character in
  51.              let asString = "\(character)"
  52.              let asInt = Int(asString)!
  53.              let asInt8 = Int8(asInt)
  54.  
  55.              return asInt8
  56.          }
  57.      }
  58.  
  59.      init(characters: [SequenceCharacter]) {
  60.         var value = 0
  61.         var p = 1
  62.  
  63.         for char in characters.reversed() {
  64.             value += p * Int(char)
  65.  
  66.             p *= 10
  67.         }
  68.  
  69.         self.init(value)
  70.      }
  71.  }
  72.  
  73.  
  74. // This defines An alphabet with the latin leters from a to z
  75. class LowercaseLetters: Alphabet {
  76.     typealias AlphabetCharacter = Character
  77.  
  78.     func allCharacters() -> [AlphabetCharacter] {
  79.         return Array("abcdefghijklmnopqrstuvwxyz".characters)
  80.     }
  81.  
  82.     required init() {
  83.  
  84.     }
  85. }
  86.  
  87. class DigitsAlphabet: Alphabet {
  88.     typealias AlphabetCharacter = Int8
  89.  
  90.     func allCharacters() -> [AlphabetCharacter] {
  91.         return Array(0...9)
  92.     }
  93.  
  94.     required init() {
  95.  
  96.     }
  97. }
  98.  
  99.  
  100. class TrieNode<T> where T:Alphabet {
  101.     typealias character = T.AlphabetCharacter
  102.  
  103.     var children: [character:TrieNode<T>] = [:]
  104.     var prefixCount: Int = 0
  105.     var isFinal:Bool = false
  106.  
  107.     init() {
  108.  
  109.     }
  110.  
  111.     ///// called from a PARENT
  112.     func createChildFor(_ character: character) -> TrieNode<T> {
  113.         let node = TrieNode<T>()
  114.  
  115.         children[character] = node
  116.  
  117.         return node
  118.     }
  119.  
  120.     func getOrCreateChildFor(_ character: character) -> TrieNode<T> {
  121.         if let child = children[character] {
  122.             return child
  123.         } else {
  124.             return createChildFor(character)
  125.         }
  126.     }
  127. }
  128.  
  129. class Trie<T> where T:Alphabet {
  130.     typealias character = T.AlphabetCharacter
  131.     typealias Node = TrieNode<T>
  132.  
  133.     var root = Node()
  134.  
  135.     func validate(characters: [character]) -> Bool {
  136.         let allCharacters = T().allCharacters()
  137.         return characters.reduce(true) { result, char in
  138.             return result && allCharacters.contains(char)
  139.         }
  140.     }
  141.  
  142.     func insert<U:CharacterSequence>(_ word: U) where U.SequenceCharacter == character {
  143.         insert(characters: word.toSequence())
  144.     }
  145.  
  146.     func insert(characters: [character]) {
  147.         guard validate(characters: characters) else {
  148.             return
  149.         }
  150.  
  151.         var node = root
  152.         node.prefixCount += 1
  153.         for character in characters {
  154.             node = node.getOrCreateChildFor(character)
  155.             node.prefixCount += 1
  156.         }
  157.         node.isFinal = true
  158.     }
  159.  
  160.     func getNodeFor(characters: [character]) -> Node? {
  161.         var node: Node? = root
  162.         for character in characters {
  163.             node = node?.children[character]  
  164.         }
  165.         return node
  166.     }
  167.  
  168.     func query<U:CharacterSequence>(_ word: U) -> Bool where U.SequenceCharacter == character {
  169.         return query(characters: word.toSequence())
  170.     }
  171.  
  172.     func query(characters: [character]) -> Bool {
  173.         guard validate(characters: characters) else {
  174.             return false
  175.         }
  176.         let node = getNodeFor(characters: characters)
  177.  
  178.         if node == nil {
  179.             return false
  180.         }
  181.         return node!.isFinal
  182.     }
  183.  
  184.  
  185.     func getNodeSequence(characters: [character]) -> [(char: character, node: Node)] {
  186.         var node: Node? = root
  187.         var nodes:[(character,Node)] = []
  188.  
  189.         let tup = (characters[0],node!)
  190.         nodes.append(tup)
  191.  
  192.         for character in characters {
  193.             node = node?.children[character]
  194.             if node == nil {
  195.                 return nodes
  196.             }
  197.             let tup = (character, node!)
  198.             nodes.append(tup)
  199.         }
  200.  
  201.         return nodes
  202.     }
  203.  
  204.  
  205.     func remove<U:CharacterSequence>(_ word: U) where U.SequenceCharacter == character {
  206.         remove(characters: word.toSequence())
  207.     }
  208.  
  209.     func remove(characters: [character]) {
  210.         guard validate(characters: characters) else {
  211.             return
  212.         }
  213.  
  214.         let charactersNodes = getNodeSequence(characters: characters)
  215.  
  216.         if charactersNodes.count != characters.count + 1 {
  217.             return
  218.         }
  219.  
  220.         var parent = root
  221.  
  222.         for (char, node) in charactersNodes {
  223.             node.prefixCount -= 1
  224.  
  225.             if node.prefixCount == 0 && node !== root {
  226.                 parent.children.removeValue(forKey: char)
  227.                 break // we deleted the reference to an entire sequence of nodes
  228.                 // swift will delete those nodes from memory
  229.             }
  230.             parent = node
  231.         }
  232.  
  233.         charactersNodes.last?.node.isFinal = false
  234.     }
  235.  
  236.  
  237.  
  238.     func removeWordsPrefix<U:CharacterSequence>(_ word: U) where U.SequenceCharacter == character {
  239.         removeWordsPrefix(characters: word.toSequence())
  240.     }
  241.  
  242.     func removeWordsPrefix(characters: [character]) {
  243.         guard validate(characters: characters) else {
  244.             return
  245.         }
  246.  
  247.         let charactersNodes = getNodeSequence(characters: characters)
  248.  
  249.         if charactersNodes.count != characters.count + 1 {
  250.             return
  251.         }
  252.  
  253.         let decreasePrefix = charactersNodes.last!.node.prefixCount
  254.  
  255.         var parent = root
  256.  
  257.         for (char, node) in charactersNodes {
  258.             node.prefixCount -= decreasePrefix
  259.             if node.prefixCount == 0 && node !== root {
  260.                 parent.children.removeValue(forKey: char)
  261.                 break // we deleted the reference to an entire sequence of nodes
  262.                     // swift will delete those nodes from memory
  263.             }
  264.             parent = node
  265.         }
  266.     }
  267.  
  268.     func kthWord<U:CharacterSequence>(atIndex index: Int, from node: Node? = nil, characters: [character] = []) -> U? where U.SequenceCharacter == character {
  269.         let node = node ?? root
  270.  
  271.         if index == 0 && node.isFinal {
  272.             return U(characters: characters)
  273.         }
  274.  
  275.         var skipped = 0
  276.         if node.isFinal {
  277.           skipped += 1
  278.         }
  279.  
  280.         for char in T().allCharacters() {
  281.             if let child = node.children[char] {
  282.                 if skipped + child.prefixCount > index {
  283.                     // find the word in the current child
  284.                     return kthWord(atIndex: index - skipped, from: child, characters: characters + [char])
  285.                 } else {
  286.                     // skip words from the current child
  287.                     skipped += child.prefixCount
  288.                 }
  289.             }
  290.  
  291.         }
  292.  
  293.         return nil
  294.     }
  295.  
  296.  
  297.     func wordsSamePrefix<U:CharacterSequence>(_ word: U) -> Int where U.SequenceCharacter == character {
  298.         return wordsSamePrefix(characters: word.toSequence())
  299.     }
  300.  
  301.     func wordsSamePrefix(characters: [character]) -> Int {
  302.         guard validate(characters: characters) else {
  303.             return 0
  304.         }
  305.  
  306.         if let node = getNodeFor(characters: characters) {
  307.             return node.prefixCount
  308.         } else {
  309.             return 0
  310.         }
  311.     }
  312.  
  313.     func suggestWords<U:CharacterSequence>(_ word: U) -> [U] where U.SequenceCharacter == character {
  314.         return suggestWords(characters: word.toSequence())
  315.     }
  316.  
  317.  
  318.  
  319.     func suggestWords<U:CharacterSequence>(characters: [character]) -> [U] where U.SequenceCharacter == character{
  320.           guard validate(characters: characters) else {
  321.               return []
  322.           }
  323.  
  324.           let lastNodeCharacters = getNodeFor(characters: characters)
  325.  
  326.           if lastNodeCharacters == nil {
  327.             return []
  328.           }
  329.           var discoveredNodes:[(node: Node, char: character, counter: Int)] = DFS(node: lastNodeCharacters!, char: characters[characters.count - 1])
  330.  
  331.          // discoveredNodes also includes the last character from characters parameter
  332.          // it can easily be removed here, since we know it's the first discovered node
  333.           discoveredNodes.remove(at: 0)
  334.  
  335.           return buildSeq(characters: characters, discoveredNodes: &discoveredNodes)
  336.       }
  337.  
  338. // parameters: node and corresponding character
  339. // counter == prefixCount; useful when building the sequence
  340.     func DFS(node: Node, char: character ) -> [(node: Node, char: character, counter: Int)] {
  341.  
  342.           let allCharacters = T().allCharacters()
  343.           var discoveredNodes:[(node: Node, char: character, counter: Int)] = []
  344.  
  345.           var stackOfNodes = Stack<(node: Node,char: character, counter: Int)>()
  346.           stackOfNodes.push((node: node, char: char, counter: node.prefixCount ))
  347.  
  348.           var nodeChar:(node: Node, char: character, counter: Int)
  349.  
  350.           while !stackOfNodes.isEmpty(){
  351.             nodeChar = stackOfNodes.pop()
  352.             discoveredNodes.append(nodeChar)
  353.  
  354.             for character in allCharacters {
  355.                 if let inter = nodeChar.node.children[character] {
  356.                     stackOfNodes.push((node:inter, char:character, counter:inter.prefixCount))
  357.                 }
  358.             }
  359.           }
  360.  
  361.           return discoveredNodes
  362.       }
  363.  
  364.  
  365. // takes the prefix, the nodes from DFS and builds words (suggestions)
  366.     func buildSeq<U:CharacterSequence>(characters: [character], discoveredNodes: inout [(node: Node, char: character, counter: Int)]) -> [U] where U.SequenceCharacter == character {
  367.       var res:[U] = []
  368.  
  369.       for i in 0..<discoveredNodes.count {
  370.         if discoveredNodes[i].node.isFinal == true {
  371.           var w:[character] = []
  372.           // counter tells us which nodes characters are to be included in the sequence
  373.           // if counter > 0 then the node's character belongs to the sequence and we decrement the counter
  374.           // otherwise, the node was used in other sequences and was "used up"
  375.           print("building seq starting with:",characters)
  376.           for j in 0...i{
  377.               if discoveredNodes[j].counter > 0 {
  378.                 print(discoveredNodes[j].char)
  379.                 w.append(discoveredNodes[j].char)
  380.  
  381.                 discoveredNodes[j].counter -= 1
  382.             }
  383.           }
  384.           w = characters + w
  385.  
  386.           res.append(U(characters: w))
  387.         }
  388.       }
  389.  
  390.       return res
  391.     }
  392.  
  393. }
  394.  
  395.  
  396.  
  397. let myTrie = Trie<LowercaseLetters>()
  398.  
  399. myTrie.insert("to")
  400. myTrie.insert("tour")
  401. myTrie.insert("tea")
  402. myTrie.insert("roll")
  403. myTrie.insert("round")
  404. myTrie.insert("route")
  405.  
  406. print (myTrie.suggestWords("ro"))
  407.  
  408. // let myTrie2 = Trie<DigitsAlphabet>()
  409. // myTrie2.insert(12);
  410. // myTrie2.insert(123);
  411. // myTrie2.insert(13);
  412. // myTrie2.insert(133);
  413. // myTrie2.insert(134);
  414. // myTrie2.insert(211);
  415. //
  416. // myTrie2.remove(12)
  417. // print(myTrie2.query(12))  // false
  418. // print(myTrie2.query(123))  // true
  419. //
  420. // myTrie2.removeWordsPrefix(12)
  421. // print(myTrie2.query(123))  // false
  422. // print(myTrie2.wordsSamePrefix(13))  // 3 : 13, 133, 134
  423. //
  424. // print (myTrie2.suggestWords(13))
RAW Paste Data