PythonでツリーのノードのK番目の祖先を求めるプログラム

n個のノード(0からn-1までの番号が振られている)からなる木を考えます。この木はparent配列によって表現され、parent[i]はノードiの親ノードを示します。木の根はノード0です。ここで、指定されたノードのk番目の祖先を求めるプログラムを作成します。該当する祖先が存在しない場合は-1を返します。

例えば、次のような木が与えられたとします。

この場合の出力は2になります。ノード6の1番目の祖先は5であり、2番目の祖先は2だからです。

解法のアプローチ

この問題は「バイナリリフティング(ダブリング)」と呼ばれる技法を使うことで効率的に解けます。kを2の冪乗の組み合わせに分解することで、親をたどる回数を大幅に削減できるのがポイントです。具体的には、以下の手順に従ってsolve()関数を定義します。引数はparent(親配列)、node(対象ノード)、k(祖先の次数)です。

• nodeが-1の場合(根を超えて祖先が存在しない場合)は、-1を返します。

• kが1の場合は、parent[node](1つ上の親)を返します。

• kとk-1のビットANDが0になる場合(kが2の冪乗である場合)は、kを半分ずつに分割して再帰的に計算します。
  solve(parent, solve(parent, node, k/2), k/2)

• それ以外の場合は、kの最上位ビット(msb = 2^(kのビット長 − 1))を取り出し、残り部分(k − msb)とmsbに分けて再帰的に処理します。
  solve(parent, solve(parent, node, k − msb), msb)

実装例

以下のPythonコードで実際の動作を確認してみましょう。

def solve(parent, node, k):
   if node == -1:
      return -1
   elif k == 1:
      return parent[node]
   elif not (k & k-1):
      return solve(parent, solve(parent, node, k >> 1), k >> 1)
   else:
      msb = 1 << (k.bit_length()-1)
      return solve(parent, solve(parent, node, k-msb), msb)

parent = [-1,0,0,1,2,2,5,5]
node = 6
k = 2
print(solve(parent, node, k))

入力

parent = [-1,0,0,1,2,2,5,5]、node = 6、k = 2

出力

2

アルゴリズムのポイント

この手法の核心は、「k番目の祖先 = (k/2)番目の祖先のさらに(k/2)番目の祖先」という性質にあります。kが2の冪乗であれば半分ずつに分割でき、そうでない場合も最上位ビットを切り出して段階的に処理することで、再帰の深さは最大でもO(log k)に収まります。単純に親をk回たどる方式(O(k))と比べ、大きなkに対しても高速に動作するのが特徴です。