Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- 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))
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement