Tuesday, November 3, 2015

Zenefits: Implement a Trie

There is a dictionary with words e.g. a, apple, about, book, beats, etc.. 
Design a data structure to support input a prefix string, output the number of words in the dictionary with the same prefix. 
e.g. Input a, return 3, because a, apple, about with the same prefix a
Input ap, only apple has the prefix ap. So return 1

Understand the problem:
The problem clearly suggests to use a trie. The trie should supports at least two operations, 
1. insert(String s);
2. countPrefix(String s)

The challenge part is the second method, i.e. to count the number of words with the same prefix. 

So the solution would be first we implement a method called getPrefix(String s) which returns the pointer of the prefix. Consider how can calculate the number of paths for a binary tree, it equals to the number of leaf nodes. Here is the same. The number of words with the same prefix equals to the number of leaf nodes. So we start from the prefix node, do a DFS recursively until we reach to a leaf node, increment the count, and return. 

Code (Java):
import java.util.*;
public class Trie {
    private TrieNode root;
    private int count = 0;
    
    public Trie() {
        root = new TrieNode();
    }
    
    // Insert a string
    public void insert(String s) {
        TrieNode p = root;
        if (s == null || s.length() == 0) {
            return;
        }
        
        for (int i = 0; i < s.length(); i++) {
            char c = s.charAt(i);
            if (p.children[c - 'a'] == null) {
                    p.children[c - 'a'] = new TrieNode();
            }
            
            if (i == s.length() - 1) {
                p.children[c - 'a'].leaf = true;
            }
            
            p = p.children[c - 'a'];
        }
    }
    
    // Get Prefix, return null if not exist
    public TrieNode getPrefix(String s) {
        if (s == null || s.length() == 0) {
            return null;
        }
        
        TrieNode p = root;
        for (int i = 0; i < s.length(); i++) {
            char c = s.charAt(i);
            
            if (p.children[c - 'a'] == null) {
                return null;
            } else {
                p = p.children[c - 'a'];
            }
        }
        
        return p;
    }
    
    // Count the number of nodes with the given prefix
    public int countPrefix(String s) {
        TrieNode p = getPrefix(s);
        count = 0;
        
        countPrefixHelper(p);
        
        return count;
    }
    
    private void countPrefixHelper(TrieNode p) {
        if (p.leaf) {
            count++;
        }
        
        for (int i = 0; i < 26; i++) {
            if (p.children[i] != null) {
                countPrefixHelper(p.children[i]);
            }
        }
    }
    
    public static void main(String[] args) {
        Trie trie = new Trie();
        
        Scanner sc = new Scanner(System.in);
        
        while (true) {
            int opCode = sc.nextInt();
            if (opCode == 1) {
                trie.insert(sc.next());
            } else if (opCode == 2) {
                int count = trie.countPrefix(sc.next());
                System.out.println(count);
                break;
            }
        }
        
        sc.close();
    }
}

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

No comments:

Post a Comment