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.

["a", "awls", "sad", "sam", "same", "sap"] from CS 61B

Example:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
// 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.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
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:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
/* 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:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
/** 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:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
/** 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:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
/** 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:

1
2
3
4
5
6
7
8
9
/** 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:

1
2
3
4
5
6
7
8
9
10
11
12
13
/** 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();
}

Comment