Trie.js 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  1. var Class = require('./Class');
  2. var each = require('./each');
  3. exports = Class({
  4. initialize: function Trie() {
  5. this.clear();
  6. },
  7. add: function(word) {
  8. var edges = this._edges;
  9. var node = this._root;
  10. this._wordsInSubtree[node]++;
  11. for (var i = 0, len = word.length; i < len; i++) {
  12. var edge = word[i];
  13. var next = edges[node][edge];
  14. if (!next) {
  15. if (this._freeNodes.length) {
  16. next = this._freeNodes.pop();
  17. } else {
  18. next = this._idx++;
  19. this._isWord.push(false);
  20. this._wordsInSubtree.push(0);
  21. edges.push({});
  22. }
  23. edges[node][edge] = next;
  24. }
  25. this._wordsInSubtree[next]++;
  26. node = next;
  27. }
  28. this._isWord[node] = true;
  29. },
  30. remove: function(word) {
  31. if (!this.has(word)) {
  32. return;
  33. }
  34. var node = this._root;
  35. this._wordsInSubtree[node]--;
  36. for (var i = 0, len = word.length; i < len; i++) {
  37. var edge = word[i];
  38. var next = this._edges[node][edge];
  39. if (!--this._wordsInSubtree[next]) {
  40. delete this._edges[node][edge];
  41. this._freeNodes.push(next);
  42. }
  43. node = next;
  44. }
  45. this._isWord[node] = false;
  46. },
  47. has: function(word) {
  48. var node = this._root;
  49. for (var i = 0, len = word.length; i < len; i++) {
  50. node = this._edges[node][word[i]];
  51. if (!node) {
  52. return false;
  53. }
  54. }
  55. return this._isWord[node];
  56. },
  57. words: function(prefix) {
  58. var node = this._root;
  59. for (var i = 0, len = prefix.length; i < len; i++) {
  60. node = this._edges[node][prefix[i]];
  61. if (!node) {
  62. return [];
  63. }
  64. }
  65. var result = [];
  66. this._dfs(node, prefix, result);
  67. return result;
  68. },
  69. clear: function() {
  70. this._idx = 1;
  71. this._root = 0;
  72. this._edges = [{}];
  73. this._isWord = [false];
  74. this._wordsInSubtree = [0];
  75. this._freeNodes = [];
  76. },
  77. _dfs: function(node, prefix, result) {
  78. var _this = this;
  79. if (this._isWord[node]) {
  80. result.push(prefix);
  81. }
  82. var edges = this._edges[node];
  83. each(edges, function(node, edge) {
  84. return _this._dfs(node, prefix + edge, result);
  85. });
  86. }
  87. });
  88. module.exports = exports;