BlacklistTrieExample.swift

// Copyright © 2019 Nicholas Peterson. All rights reserved.
//

import Foundation

protocol BlacklistInterface {
    /*
     * Creates a new list from an array of blacklisted words.
     * Runs at worst in O(n * m) time where n is the length of the longest word and m is the total number of words.
     */
    static func fromList(_ wordList: [String]) -> BlacklistInterface
    
    /*
     * Adds word to Blacklist
     * If a shorter root already exists nothing will be added.
     * If the new word is a prefix to other words, the existing longer words are removed, and the new shorter prefix is added.
     * Runs in O(n) time (at worst) when n is the length of the word being added.
     */
    func addWord(_ word: String)

    /*
     * Return true if a given word, or prefix of a give word, exists in the blacklist.
     * Runs in O(n) time where n is the length of the word.
     */
    func isBlacklisted(_ word: String) -> Bool
    
    /*
     * Returns an array of words in our blacklist.
     * Runs commonly in O(n log(m)) n is the number of words, and m is the longest stored word.
     * In the worst case where all words are unique and the same length, time becomes O(n * m).
     */
    func words() -> [String]

}

// A Trie structure that represents our blacklist
class Blacklist: BlacklistInterface {
    private var isWord : Bool { return !isHead && characterMap.isEmpty }
    private var isHead : Bool { return value == "" }
    private let value: String
    private var characterMap = [Character: Blacklist]()
    
    init(value: String = "") { // Empty string indicates head/root
        self.value = value
    }
    
    class func fromList(_ wordList: [String]) -> BlacklistInterface {
        let blacklist = Blacklist()
        wordList.forEach { blacklist.addWord($0.lowercased()) }
        
        return blacklist
    }
    
    func addWord(_ word: String) {
        internal_addWord(word)
    }

    func isBlacklisted(_ word: String) -> Bool {
        return seekWord(word).isWord
    }
    
    func words() -> [String] {
        guard !isWord else { return [value] }
        
        return characterMap.values.flatMap { $0.words() }
    }
    
    
    // Return a new Blacklist representation of a single word.
    // A partial word will be return if `stoppingAt` is given a value the includes the first common character.
    // Runs in O(n) time where n is the length of the word.
    private class func newWord(_ word: String, to existingPart: Blacklist? = nil, stoppingAt stopWord: String? = nil) -> Blacklist {
        let previousPart = existingPart ?? Blacklist(value: word)
        guard !word.isEmpty, word != stopWord  else { return previousPart }
        
        let trimmedWord = String(word[..<word.lastCharacterIndex])
        let lastCharacter = word[word.lastCharacterIndex]
        let newPart = Blacklist(value: trimmedWord)
        newPart.characterMap[lastCharacter] = previousPart
        
        return newWord(trimmedWord, to: newPart, stoppingAt: stopWord)
        
    }
    
    private func internal_addWord(_ word: String, to existingPart: Blacklist? = nil) {
        // Find our intersection.
        let existingBlacklist = existingPart ?? seekWord(word)
        if existingBlacklist.value == word || existingBlacklist.isWord {
            existingBlacklist.characterMap = [:]
            return
        }
        
        // Work backwards to our intersection.
        let remainder = Blacklist.newWord(word, stoppingAt: existingBlacklist.value)
        
        // ... and link them up
        // We assume our new words first element is always the same as self.
        remainder.characterMap.keys.forEach { c in // Based on the above, this will only one once.
            existingBlacklist.characterMap[c] = remainder.characterMap[c]
        }
    }
    
    // Get the intersection between the word and the dictionary
    // Runs in O(n) time where n is the length of the word.
    private func seekWord(_ word: String, afterIndex after: String.Index? = nil) -> Blacklist {
        if isWord { return self }
        let i = after ?? word.startIndex
        let wordSpace = word.startIndex..<word.endIndex
        
        guard wordSpace.contains(i) else { return self }
        let c = word[i]
        
        if let existing = characterMap[c] {
            return existing.seekWord(word, afterIndex: word.index(after: i))
        }
        return self
    }
}

extension Blacklist: CustomDebugStringConvertible {
    var debugDescription: String {
        return "[\(words().joined(separator: ", "))]"
    }
}

fileprivate extension String {
    var lastCharacterIndex: String.Index {
        guard endIndex != startIndex else { return startIndex }
        return index(endIndex, offsetBy: -1)
    }
}