C++で解説：T秒後のカエルの位置を求める確率計算アルゴリズム

n個の頂点からなる無向木(ツリー)があるとします。頂点には1からnまでの番号が付けられており、カエルは頂点1からジャンプを開始します。カエルは、現在いる頂点に隣接している「未訪問」の頂点へ、1秒でジャンプすることができますが、一度訪れた頂点へ戻ることはできません。ジャンプ先の候補が複数ある場合は、いずれも等しい確率でランダムに1つを選んで移動します。逆に、行ける未訪問の頂点がなくなったカエルは、その場で永遠に跳ね続けることになります。

木は辺の配列として与えられます。ここで求めたいのは、「t秒後にカエルが頂点targetの上にいる確率」です。

問題の例

たとえば、入力が n = 7、t = 2、target = 4 で、木が次図のような形状だったとしましょう。

このとき出力は 0.1666 になります。カエルは頂点1からスタートし、1秒後に確率 1/3(≒0.3333)で頂点2へジャンプし、さらに2秒後に確率 1/2(=0.5)で頂点4へジャンプするからです。したがって、2秒後にカエルが頂点4にいる確率は 0.3333 × 0.5 = 約 0.16666(≒ 1/6)となります。

解法の考え方

この問題は、頂点1からtargetまでの経路を辿りながら、各ステップのジャンプ確率を掛け合わせることで解けます。ポイントは、target側から頂点1へ向かってDFS(深さ優先探索)を行い、経路と所要時間を記録する点です。各頂点での遷移確率は「その時点で選べる隣接頂点の数」で決まり、根(頂点1)以外では来た方向へ戻れないため、分母は「隣接頂点数 − 1」になります。

アルゴリズムの手順

• ret := 1 で初期化し、訪問済み頂点を管理する集合 visited を用意します。
• DFS部分： 関数 dfs(node, start, 辺リスト g, time, t, スタック st) を定義します。
  • node がすでに visited に含まれている場合は false を返します。
  • node を visited に追加します。
  • node が頂点1に到達したら、tt := time(所要時間)、ok := true とし、true を返します。
  • i := 0 から g[node] のサイズ未満の間、次を繰り返します。
    • g[node][i] をスタック st に積みます。
    • dfs(g[node][i], start, g, time + 1, t, st) が true を返したら、そのまま true を返します。
    • 失敗したら st から要素を取り除きます(バックトラック)。
  • 最後に false を返します。
• メイン処理:
  • ret := 1、ok := false で初期化します。
  • サイズ n+1 の隣接リスト graph、graph2 を用意し、edges の各辺を双方向に登録します。
  • 空のスタック st を用意し、dfs(target, target, graph, 0, t, st) を呼び出します。
  • st が空になるまで、次を繰り返します。
    • node := st のトップ要素、sz := graph[node] のサイズとします。
    • node が頂点1でない場合は、sz を1減らします。
    • ret := ret × (1.0 / sz) として確率を掛け合わせ、st からポップします。
  • tt > t なら 0 を返します(targetへの到達に t 秒より多くかかるケース)。
  • tt == t なら、ちょうど t 秒時に target へ到達するので ret を返します。
  • tt < t かつ target == 1 かつ graph[target].size() ≥ 1 の場合、カエルが直ちに頂点1を出て行ってしまうため 0 を返します。
  • それ以外は、tt < t かつ graph[target].size() > 1(t秒より早く着き、なおかつ葉ではないので途中で飛び出してしまう)なら 0、そうでなければ(葉にとどまり続けるため)ret を返します。

それでは、実際の実装を見て理解を深めましょう。

C++での実装例

#include <bits/stdc++.h>
using namespace std;
class Solution {
   public:
   double ret = 1;
   bool ok;
   set<int> visited;
   int tt;
   bool dfs(int node, int start, vector<int> g[], int time, int t,
   stack<int>& st){
      if (visited.count(node))
      return false;
      visited.insert(node);
      if (node == 1) {
         tt = time;
         ok = true;
         return true;
      }
      for (int i = 0; i < g[node].size(); i++) {
         st.push(g[node][i]);
         if (dfs(g[node][i], start, g, time + 1, t, st))
         return true;
         ;
         st.pop();
      }
      return false;
   }
   double frogPosition(int n, vector<vector<int> >& edges, int t,
   int target){
      ret = 1;
      ok = false;
      vector<int> graph[n + 1];
      vector<int> graph2[n + 1];
      for (int i = 0; i < edges.size(); i++) {
         graph[edges[i][0]].push_back(edges[i][1]);
         graph[edges[i][1]].push_back(edges[i][0]);
      }
      stack<int> st;
      dfs(target, target, graph, 0, t, st);
      while (!st.empty()) {
         int node = st.top();
         double sz = (double)graph[node].size();
         if (node != 1)
         sz--;
         ret *= (1.0 / sz);
         st.pop();
      }
      if (tt > t)
      return 0;
      if (tt == t)
      return ret;
      if (tt < t && target == 1 && graph[target].size() >= 1)
      return 0;
      return tt < t && graph[target].size() > 1 ? 0 : ret;
   }
};
main(){
   Solution ob;
   vector<vector<int>> v = {{1,2},{1,3},{1,7},{2,4},{2,6},{3,5}};
   cout << (ob.frogPosition(7,v,2,4));
}

入力

7, {{1,2},{1,3},{1,7},{2,4},{2,6},{3,5}}, 2, 4

出力

0.166667

まとめ

このアルゴリズムは、targetから頂点1へ向かう経路を1回のDFSで検出し、経路上の各頂点における遷移確率を掛け合わせることで答えを得ます。木の全頂点と辺を高々1度ずつ処理するため、計算量は頂点数を n とすると O(n) と非常に効率的です。「ちょうどt秒で到達する場合」「それより早く到達して葉にとどまる場合」「途中で飛び出してしまう場合」といった条件分岐を正しく扱うことが、この種の確率問題を解く際の重要なポイントになります。