Implement Trie (Prefix Tree)

208. Implement Trie (Prefix Tree)

LintCode-442.Implement Trie

Implement a trie with insert, search, and startsWith methods.

Note:
You may assume that all inputs are consist of lowercase letters a-z.

class TrieNode {
    private var children: [TrieNode?] = []
    public var isEnd = false

    init() {
        self.children = Array(repeating: nil, count: 26)
        self.isEnd = false
    }

    func put(_ c: Character, _ node: TrieNode) {
        let val = Int(c.unicodeScalars.first!.value - UnicodeScalar("a").value)
        children[val] = node
    }

    func get(_ c: Character) -> TrieNode? {
        return children[Int(c.unicodeScalars.first!.value - UnicodeScalar("a").value)]
    }
}

class Trie {
    private var root: TrieNode

    /** Initialize your data structure here. */
    init() {
        self.root = TrieNode()
    }
    
    /** Inserts a word into the trie. */
    func insert(_ word: String) {
        var node = root
        
        for c in word {
            if let n = node.get(c) {
                node = n
            } else {
                let n = TrieNode()
                node.put(c, n)
                
                node = n
            }
        }
        
        node.isEnd = true
    }
    
    func searchPrefix(_ word: String) -> TrieNode? {
        var node = root

        for c in word {
            if let n = node.get(c) {
                node = n
            } else {
                return nil
            }
        }

        return node
  }
    
    /** Returns if the word is in the trie. */
    func search(_ word: String) -> Bool {
        let node = searchPrefix(word)
        
        return node?.isEnd ?? false
    }
    
    /** Returns if there is any word in the trie that starts with the given prefix. */
    func startsWith(_ prefix: String) -> Bool {
        let node = searchPrefix(prefix)
        
        return node != nil
    }
}

/**
 * Your Trie object will be instantiated and called as such:
 * let obj = Trie()
 * obj.insert(word)
 * let ret_2: Bool = obj.search(word)
 * let ret_3: Bool = obj.startsWith(prefix)
 */

class TrieNode {
    private TrieNode[] children;
    private boolean isEnd;

    public TrieNode() {
        children = new TrieNode[26];
    }

    public boolean containsKey(char c) {
        return children[c -'a'] != null;
    }
    
    public TrieNode get(char c) {
        return children[c -'a'];
    }
    
    public void put(char c, TrieNode node) {
        children[c -'a'] = node;
    }
    
    public void setEnd() {
        isEnd = true;
    }
    
    public boolean isEnd() {
        return isEnd;
    }
}

public class Trie {
    private TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    // Inserts a word into the trie.
    public void insert(String word) {
        TrieNode node = root;
        
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            
            if (node.containsKey(c)) {
                node = node.get(c);
            } else {
                TrieNode newNode = new TrieNode();
                node.put(c, newNode);
                node = newNode;
            }
        }
        
        node.setEnd();
    }
    
    // search a prefix or whole key in trie and
    // returns the node where search ends
    private TrieNode searchPrefix(String word) {
        TrieNode node = root;
        
        for (int i = 0; i < word.length(); i++) {
           char c = word.charAt(i);
           
           if (node.containsKey(c)) {
               node = node.get(c);
           } else {
               return null;
           }
        }
        
        return node;
    }

    // Returns if the word is in the trie.
    public boolean search(String word) {
        TrieNode node = searchPrefix(word);
        
        return node != null && node.isEnd();
    }

    // Returns if there is any word in the trie
    // that starts with the given prefix.
    public boolean startsWith(String prefix) {
        TrieNode node = searchPrefix(prefix);
        
        return node != null;
    }
}

// Your Trie object will be instantiated and called as such:
// Trie trie = new Trie();
// trie.insert("somestring");
// trie.search("key");