208. Implement Trie (Prefix Tree)

Reference: LeetCode
Difficulty: Medium

My Post: Java Solution for Required and Extra Methods (collect, longestPrefixOf, etc)

Problem

Implement a trie with insert, search, collect, and startsWith, wordsWithPrefix, longestPrefixOf methods.

Example:

// LeetCode Example
Trie trie = new Trie();

trie.insert("apple");
trie.search("apple");          // returns true
trie.search("app");            // returns false
trie.startsWith("app");        // returns true
trie.insert("app");   
trie.search("app");            // returns true
trie.collect();                // ["app", "apple"]
trie.wordsWithPrefix("");      // ["app", "apple"]
trie.wordsWithPrefix("appl");  // ["apple"]
trie.wordsWithPrefix("appx");  // []
trie.longestPrefixOf("appleiscool"); // "apple"

Analysis

Reference: CS 61B | Part 8 | Tries (Prefix), Range Searching and Multi-Dimensional Data

Check out the reference for details.

class Trie { 
  // Node class
  private class Node {
    Node[] next;
    boolean isKey; // indicate if it is an ending node
    Node(int r, boolean k) {
      next = new Node[r]; // R is 26 in this case
      isKey = k;
    }
  }
  
  // Members
  private static final int R = 26;
  private Node root; // act as a dummy
  
  public Trie() {
    root = new Node(R, false);
  }

  public void insert(String s) {}
  public boolean search(String s) {}
  public List<String> collect() {}
  private void collectHelper(Node p, StringBuilder sb, List<String> result) {}
  public boolean startsWith(String prefix) {} // similar to the code in search
  private Node startsWithHelper(String prefix) {} // return the last node if prefix matches.
  public List<String> wordsWithPrefix(String prefix) {}
  public String longestPrefixOf(String s) {}
}

insert:

/* Inserts a word into the trie. */
public void insert(String s) {
  if (s == null || s.length() == 0) return;
  int n = s.length();
  Node p = root;
  for (int i = 0; i < n; ++i) {
    char ch = s.charAt(i);
    if (ch < 'a' || ch > 'z') throw new IllegalArgumentException("not a-z");
    if (p.next[ch - 'a'] == null) {
      p.next[ch - 'a'] = new Node(R, false); // false as default
    }
    p = p.next[ch - 'a']; // go to the current node
    if (i == n - 1) { // last one
      p.isKey = true;
    }
  }

search:

/** Contains a word */
public boolean search(String s) {
  if (s == null || s.length() == 0) return true;
  int n = s.length();
  Node p = root;
  for (int i = 0; i < n; ++i) {
    char ch = s.charAt(i); // assume ch is ['a', 'z']
    if (p.next[ch - 'a'] == null) return false;
    p = p.next[ch - 'a'];
    if (i == n - 1) { // last one
      if (p.isKey == false) return false;
    }
  }
  return true;
}

collect:

/** Returns all the words is in the trie. */
public List<String> collect() {
  List<String> result = new ArrayList<>();
  StringBuilder sb = new StringBuilder();
  collectHelper(root, sb, result);
  return result;
}

private void collectHelper(Node p, StringBuilder sb, List<String> result) {
  if (p == null) return;
  if (p.isKey) {
    result.add(sb.toString());
    // do not return yet
  }
  for (int i = 0; i < p.next.length; ++i) {
    Node nextNode = p.next[i];
    char ch = (char) ('a' + i);
    sb.append(ch);
    collectHelper(nextNode, sb, result);
    sb.setLength(sb.length() - 1);
  }
}

startsWith:

/** Returns if there is any word in the trie that starts with the given prefix. */
// similar to the code in search
public boolean startsWith(String prefix) {
  if (prefix == null || prefix.length() == 0) return true;
  return startsWithHelper(prefix) != null;
}

// return the last node if prefix matches.
private Node startsWithHelper(String prefix) {
  int n = prefix.length();
  Node p = root;
  for (int i = 0; i < n; ++i) {
    char ch = prefix.charAt(i); // assume ch is ['a', 'z']
    if (p.next[ch - 'a'] == null) return null;
    p = p.next[ch - 'a'];
    // don't care if the last node is key or not
  }
  return p;
}

wordsWithPrefix:

/** Returns all words that start with the given prefix. */
public List<String> wordsWithPrefix(String prefix) {
  // find the node corresponding with the last char in prefix
  List<String> result = new ArrayList<>();
  Node p = startsWithHelper(prefix);
  StringBuilder sb = new StringBuilder(prefix);
  collectHelper(p, sb, result);
  return result;
}

longestPrefixOf:

/** Returns the longest prefix of the word s. */
public String longestPrefixOf(String s) {
  int n = s.length;
  StringBuilder sb = new StringBuilder(); // construct the prefix
  Node p = root;
  for (int i = 0; i < n; ++i) {
    char ch = s.charAt(i); // assume ch is ['a', 'z']
    if (p.next[ch - 'a'] == null) break;
    p = p.next[ch - 'a'];
    sb.append(ch);
  }
  return sb.toString();
}