【C++】パス合計IV（Path Sum IV）― 根から葉までの経路の合計をDFSで求める

問題の概要

深さが5より小さい二分木を、3桁の整数のリストで表すことを考えます。木の深さが5未満であれば、この木は3桁の整数のリストとして完全に表現できます。リスト内の各整数は、次のような意味を持っています。

• 百の位：そのノードの深さ D（1 ≤ D ≤ 4）を表します。

• 十の位：そのノードが属するレベル内での位置 P（1〜8）を表します。位置の付け方は完全二分木の場合と同じです。

• 一の位：そのノードの値 V（0 ≤ V ≤ 9）を表します。

求めたいのは、根（ルート）から葉までのすべての経路における値の合計です。

たとえば入力が [113, 215, 221] の場合、出力は 12 になります。このリストが表す木は次の図のとおりです。

経路の合計は (3 + 5) + (3 + 1) = 12 となります。

解き方（アルゴリズム）

この問題は、深さ優先探索（DFS）を使うことで効率よく解けます。手順は以下のとおりです。

• マップ型のコンテナ graph を定義します。

• 関数 dfs() を定義します。引数は node、level、pos、sum（初期値 0）です。

• isLeaf を true に初期化します。

• i を 0 から graph[level + 1] のサイズ未満まで増やしながら、以下を繰り返します。

  • ペア temp := graph[level + 1][i] を取り出します。

  • temp.first / 2 が pos と一致する場合：

    • isLeaf := false とします。

    • dfs(temp.second, level + 1, temp.first, sum + node) を再帰的に呼び出します。

• isLeaf が true のまま（＝葉に到達した）場合：

  • ret := ret + (sum + node) として答えに加算します。

メインメソッドでは、以下の処理を行います。

• ret := 0 で初期化します。

• i を 0 から nums のサイズ未満まで増やしながら、以下を繰り返します。

  • x := nums[i]

  • val := x mod 10（一の位＝ノードの値）

  • x := x / 10

  • pos := x mod 10（十の位＝レベル内の位置）

  • x := x / 10

  • level := x（百の位＝深さ）

  • { 1 を (level − 1) 回左シフトした値 ＋ pos − 1, val } を graph[level] の末尾に追加します。

• dfs(graph[1][0].second, 1, graph[1][0].first) を呼び出します。

• ret を返します。

仕組みのポイント

このアルゴリズムが正しく動く理由は、完全二分木の性質にあります。深さ d・位置 p のノードの子は、次のレベルの位置 2p−1 と 2p に存在します。つまり「子の位置を 2 で割ると親の位置になる」ため、temp.first / 2 == pos という比較だけで親子関係を判定できます。また、(1 << (level − 1)) + pos − 1 というキーにより、各ノードを木全体の中で一意な番号に変換しています。DFS で葉に到達した時点で、それまでの累積和 sum ＋ 現在のノードの値を ret に足し合わせることで、すべての根〜葉経路の合計が得られます。

C++による実装例

理解を深めるために、実際の C++ コードを見てみましょう。

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
   int ret;
   map <int, vector < pair <int, int> > > graph;
   void dfs(int node, int level, int pos, int sum = 0){
      bool isLeaf = true;
      for (int i = 0; i < graph[level + 1].size(); i++) {
         pair<int, int> temp = graph[level + 1][i];
         if (temp.first / 2 == pos) {
            isLeaf = false;
            dfs(temp.second, level + 1, temp.first, sum + node);
         }
      }
      if (isLeaf) {
         ret += (sum + node);
      }
   }
   int pathSum(vector<int>& nums) {
      ret = 0;
      for (int i = 0; i < nums.size(); i++) {
         int x = nums[i];
         int val = x % 10;
         x /= 10;
         int pos = x % 10;
         x /= 10;
         int level = x;
         graph[level].push_back({ (1 << (level - 1)) + pos - 1, val });
      }
      dfs(graph[1][0].second, 1, graph[1][0].first);
      return ret;
   }
};
main(){
   Solution ob;
   vector<int> v = {113,215,221};
   cout<<(ob.pathSum(v));
}

入力

{113,215,221}

出力

12