Python で非再帰 DFS を楽に実装したい話

ABC218 で非常にお世話になったので...

実装概要

• (返り値・再帰のときの引数にあたるものを記録しておくための配列を作る)
• 各頂点で処理をするとき、
  • はじめに [現在処理をしている頂点, 0] を deque に append する
  • その後、移動可能頂点があるならば、 [移動先の頂点, 1] を順番に deque に append していく
• (返り値がある場合、que.pop()[1] が 0 ならば移動した頂点から返り値を収集し(配列に記録してある)、処理を行う)

こうすることで、順番に pop するだけで再帰と同じ操作を行うことができる。
言葉だけでは分かり辛いため、以下のような木を探索することを考える。 頂点番号が探索順となるように探索してみると、deque の中身は

[[1, 1]]
#頂点1の処理
[[1, 0], [5, 1], [2, 1]]
#頂点2の処理
[[1, 0], [5, 1], [2, 0], [4, 1], [3, 1]]
#頂点3の処理&後処理
[[1, 0], [5, 1], [2, 0], [4, 1]]
#頂点4の処理&後処理
[[1, 0], [5, 1], [2, 0]]
#頂点2の後処理・頂点5の処理
[[1, 0], [5, 0], [6, 1]]
#頂点6の処理&後処理・頂点5の後処理・頂点1の後処理
[]

というようになる。

実装

この実装方法を用いて、試しに ABC165 F - LIS on Tree を解いてみる。

コード例 (PyPy3, 672ms)(提出先)

from collections import deque

n = int(input())
a = list(map(int,input().split()))
path = [[] for _ in range(n)]
for i in range(n-1):
    u, v = map(int,input().split())
    path[u-1].append(v-1)
    path[v-1].append(u-1)

lis = [10**10] * n

def update(x):
    ok = n
    ng = 0
    while abs(ok-ng) > 1:
        ind = (ok + ng) // 2
        if lis[ind-1] >= x:
            ok = ind
        else:
            ng = ind
    return ok

que = deque()
que.append([0, 1])
seen = [False] * n
seen[0] = True
leng = 0
ans = [0] * n
update_memo = [[0] * 3 for _ in range(n)]

while len(que):
    task = que.pop()
    node = task[0]
    if task[1]:
        upd = update(a[node])
        update_memo[node] = [upd-1, lis[upd-1], leng]
        leng = max(leng, upd)
        lis[upd-1] = a[node]
        ans[node] = leng
        que.append([node, 0])
        for i in path[node]:
            if not seen[i]:
                seen[i] = True
                que.append([i, 1])
    else:
        lis[update_memo[node][0]] = update_memo[node][1]
        leng = update_memo[node][2]

for i in ans:
    print(i)

再帰関数を用いたところ、(私の実装方法では) TLE したため、明らかに高速になっていることが分かる。