SHARE
TWEET

Untitled

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