Add Trie

[moonjs0113]

Trie

Description

This is Trie data structure implementation in Swift.
You can check the comments in file changed.

Feature

• insert(word: String): Insert data.
  • insert(words: [String]): Insert multiple data.
• delete(word: String) -> Bool: Delete data.
  • delete(words: [String]) -> [String: Bool]: Delete multiple data.
• contains(word: String) -> Bool: Find the word in trie.
• getAllWords() -> [String]: Get All Words in trie.

Code

final class Trie {
    private var child: [String: Trie] = [:]
    private var isLeaf: Bool = false
    
    public init() {
        
    }

    public func insert(word: String) {
        var node = self
        for s in word {
            let key = "\(s)"
            node.child[key] = (node.child[key] ?? Trie())
            node = node.child[key]!
        }
        node.isLeaf = true
    }
    
    public func insert(words: [String]) {
        for word in words {
            self.insert(word: word)
        }
    }
    
    public func delete(word: String) -> Bool {
        var root = self
        var keys = Array(word).map(String.init)
        func _delete(_ node: Trie, _ depth: Int) -> Bool {
            if depth == word.count {
                guard node.isLeaf else {
                    return false
                }
                node.isLeaf = false
                return node.child.isEmpty
            }
            guard let c = node.child[keys[depth]] else {
                return false
            }
            let isDeleteChildNode = _delete(c, depth + 1)
            
            if isDeleteChildNode {
                if c.child.isEmpty && !c.isLeaf {
                    node.child[keys[depth]] = nil
                }
            }
            return isDeleteChildNode
        }
        return _delete(root, 0)
    }
    
    public func delete(words: [String]) -> [String: Bool] {
        var result: [String: Bool] = [:]
        for word in words {
            result[word] = self.delete(word: word)
        }
        return result
    }
    
    public func contains(word: String) -> Bool {
        var node = self
        for s in word {
            guard let c = node.child["\(s)"] else {
                return false
            }
            node = c
        }
        return node.isLeaf
    }
    
    public func getAllWords() -> [String] {
        var node = self
        var result: [String] = []
        for c in node.child {
            if c.value.isLeaf {
                result += ["\(c.key)"]
            }
            result += c.value.getAllWords().map { "\(c.key)\($0)" }
        }
        return result
    }
}

Test Code

let trie = Trie()
trie.insert(words: ["Swift", "Apple", "iPhone", "iPad", "iMac", "AppleWatch"])
print(trie.getAllWords())

print(trie.contains(word: "Apple"))
print(trie.delete(words: ["Apple", "iP"]))
print(trie.contains(word: "Apple"))
print(trie.getAllWords())

print(trie.contains(word: "AppleWatch"))
print(trie.contains(word: "iP"))