实现 Trie (前缀树)

[louzhedong]

习题

出处 LeetCode 算法第208题

实现一个 Trie (前缀树)，包含 insert, search, 和 startsWith 这三个操作。

示例:

Trie trie = new Trie();

trie.insert("apple");
trie.search("apple"); // 返回 true
trie.search("app"); // 返回 false
trie.startsWith("app"); // 返回 true
trie.insert("app");
trie.search("app"); // 返回 true
说明:

你可以假设所有的输入都是由小写字母 a-z 构成的。
保证所有输入均为非空字符串。

思路

用map来保存当前字母后续的字母，并用一个变量来保存当前字母是否为一个单词的末尾

解答

JavaScript

/**
 * Initialize your data structure here.
 */
var TrieNode = function () {
  this.links = new Map();
}

var Trie = function () {
  this.root = new TrieNode();
  this.root.links.set(0, true);
};

/**
 * Inserts a word into the trie. 
 * @param {string} word
 * @return {void}
 */
Trie.prototype.insert = function (word) {
  var node = this.root;
  for (var w of word) {
    if (node.links.has(w)) {
      node = node.links.get(w);
    } else {
      node.links.set(w, new TrieNode());
      node = node.links.get(w);
    }
  }
  node.links.set(0, true);
};

/**
 * Returns if the word is in the trie. 
 * @param {string} word
 * @return {boolean}
 */
Trie.prototype.search = function (word) {
  var node = this.root;
  for (var w of word) {
    if (!node.links.has(w)) {
      return false;
    }
    node = node.links.get(w);
  }
  return node.links.has(0);
};

/**
 * Returns if there is any word in the trie that starts with the given prefix. 
 * @param {string} prefix
 * @return {boolean}
 */
Trie.prototype.startsWith = function (prefix) {
  var node = this.root;
  for (var w of prefix) {
    if (!node.links.has(w)) {
      return false;
    }
    node = node.links.get(w);
  }
  return true;
};