Implement Trie (Prefix Tree)

A trie use a data structure as follows:

structure Node
   Children Node[Alphabet-Size]
   Is-Terminal Boolean
   Value Data-Type
end structure

It's a recursive data structure, a Node has an attribute Children as a Node map. Usually, when we save a list of words in a trie, the key of the Node map is one character of word and the value of map is a Node which contains next character in the word.

Say we want to save three word ["ab", "ac", "abd"] in a trie, then the data structure is like follows:

{
    children: {
        a: {
            children: {
                b: {
                    children: {
                        d: {
                            children: {},
                            isWord: true
                        }
                    },
                    isWord: false
                },
                c: {
                    children: {}
                    isWord: true
                }
            },
            isWord: fasle
        }
    }
    isWord: false
}

Let's implement it in code.

class Trie {
    private m: Map<string, Trie>;
    private isWord: boolean;

    constructor() {
        this.m = new Map();
        this.isWord = false;
    }

    insert(word: string): void {
        let cur: Trie = this;
        for (const char of word) {
            let next = cur.m.get(char);
            if (!next) {
                next = new Trie();
                cur.m.set(char, next);
            }
            cur = next;
        }
        cur.isWord = true;
    }

    search(word: string): boolean {
        let cur: Trie = this;
        for (const char of word) {
            const next = cur.m.get(char);
            if (!next) {
                return false;
            }
            cur = next;
        }
        return cur.isWord;
    }

    startsWith(prefix: string): boolean {
        let cur: Trie = this;
        for (const char of prefix) {
            const next = cur.m.get(char);
            if (!next) {
                return false;
            }
            cur = next;
        }
        return true;
    }
}

Word Search II

Now we have implemented the trie data structure, let's try to solve a leetcode problem using trie.

The problem is to find if words exist in a two dimensional board. So the process is to find all possible words in the board and then check each word to see if its in the target word set. We can use a depth first approach to find all possible words, and use a trie to check if the word is valid.

// To make word check more effectively, here some changes
// 1. only provides insert method
// 2. expose the next map and store the valid word
class Trie {
    public next: Map<string, Trie>;
    public word: string | undefined;

    constructor() {
        this.next = new Map();
    }

    insert(word: string): void {
        let cur: Trie = this;
        for (const char of word) {
            let next = cur.next.get(char);
            if (next === undefined) {
                next = new Trie();
                cur.next.set(char, next);
            }
            cur = next;
        }
        cur.word = word;
    }
}

const options = [
    [0, 1], [0, -1],
    [1, 0], [-1, 0]
];

function dfs(board: string[][], trie: Trie, x: number, y: number, res: string[]) {
    // note here
    // trie follows the dfs process
    const c = board[x][y];
    const next = trie.next.get(c);
    if (!next) return;
    trie = next;
    if (trie.word) {
        res.push(trie.word);
        trie.word = "";
    }

    board[x][y] = "#";
    for (const option of options) {
        const i = x + option[0];
        const j = y + option[1];
        if (i < 0 || j < 0) continue;
        if (i >= board.length || j >= board[x].length) continue;
        if (board[i][j] === "#") continue;
        dfs(board, trie, i, j, res);
    }
    board[x][y] = c;
}

function findWords(board: string[][], words: string[]): string[] {
    const trie = new Trie();
    for (const word of words) {
        trie.insert(word);
    }

    const res: string[] = [];
    for (let x = 0; x < board.length; x++) {
        for (let y = 0; y < board[x].length; y++) {
            dfs(board, trie, x, y, res);
        }
    }
    return res;
};