JavaScriptでMapSumクラスを実装する：トライ木を使った接頭辞の合計計算

問題

本記事では、insert メソッドと sum メソッドを持つ MapSum クラスを実装していきます。

• insert(key, val): キー（文字列）と値（整数）のペアを受け取ります。同じキーがすでに存在する場合は、既存のキーと値のペアが新しい値で上書きされます。
• sum(prefix): 接頭辞（プレフィックス）を表す文字列を受け取り、その接頭辞で始まるすべてのキーに対応する値の合計を返します。

解決のアプローチ：トライ木（接頭辞木）

この問題はトライ木（Trie／接頭辞木）というデータ構造を使うことで効率的に解くことができます。各ノードは1文字を保持し、子ノードへの参照をオブジェクト（マップ）で管理します。ノードが持つ num プロパティには、その位置で終わるキーに対応する値を格納します。

• insert の処理: ルートから順にキーの各文字をたどり、該当する子ノードがなければ新しく作成しながら進みます。そして最終的なノードに値をセットします。計算量は O(k)（k はキーの長さ）です。
• sum の処理: まず接頭辞に対応するノードまでたどります。途中で該当する子ノードが存在しなければ 0 を即座に返します。ノードが見つかったら、そこを起点として部分木全体を再帰的に走査し、すべての num の合計を算出します。

実装例

以下が実際のコードです。

class Node {
   constructor(val) {
      this.num = 0
      this.val = val
      this.children = {}
   }
}
class MapSum {
   constructor(){
      this.root = new Node('');
   }
}
MapSum.prototype.insert = function (key, val) {
   let node = this.root
   for (const char of key) {
      if (!node.children[char]) {
         node.children[char] = new Node(char)
      }
      node = node.children[char]
   }
   node.num = val
}

MapSum.prototype.sum = function (prefix) {
   let sum = 0
   let node = this.root
   for (const char of prefix) {
      if (!node.children[char]) {
         return 0
      }
      node = node.children[char]
   }
   const helper = (node) => {
      sum += node.num
      const { children } = node
      Object.keys(children).forEach((key) => {
         helper(children[key])
      })
   }
   helper(node)
   return sum
}
const m = new MapSum();
console.log(m.insert('apple', 3));
console.log(m.sum('ap'));

出力

undefined
3

コードのポイント

m.insert('apple', 3) は明示的な戻り値を持たないため、コンソールには undefined と表示されます。続く m.sum('ap') は、キー「apple」が接頭辞「ap」で始まるため、その値である 3 を返します。

さらに別のキーを追加した場合の動作も確認してみましょう。

m.insert('app', 2);
console.log(m.sum('ap')); // 5（'apple' の 3 と 'app' の 2 の合計）

このように、トライ木を活用することで、キーの登録も接頭辞による検索・合計計算もどちらも高速かつシンプルに実現できます。単語の自動補完や検索サジェストなど、接頭辞ベースの処理が必要な場面で特に有効なアプローチです。