ABC148-F Playing Tag on Treeを解いた - 少ない学びをせめて記録する

ABC148に出た。全完だったが簡単だったようでパフォーマンスは大したことなかった。

木の上で高橋くん(以下T)と青木くん(以下A)が鬼ごっこする問題、高橋くんから順に互いに1ずつ移動する。Aが鬼。最も長く移動した場合のAの移動距離を求める。基本的にAの初期位置からTの移動可能な中で最も遠いノード(Nt)までの距離-1が正解になる。

Ntまでの距離を-1する場合としない場合を偶奇などで分けるべきか本番中迷った。 AはA1,A2,,,Ak-1, Ak(=Nt)と移動すると考えると、偶奇によって両者の行動は

T:T1
A:A1
...
T : Ak
A : Ak-1
T: Ak-1 → 終了

となる場合と

T:T1
A:A1
...
T : Ak-1
A:Ak-1 →終了

となる場合があるが、回答はAの移動距離なのでどちらも同じになる。まずTから木を探索して各ノードのTからの距離を計算。その後Aからもう一度探索、ここで(Tからの距離>=Aからの距離)となっていてノードはTは到達できない。そのようなノードに到達した場合、探索における次ノードは子孫にTがいるノードとなる。

それと本番中はpythonの再帰を書かないようにしようとして詰まり結局再帰で出した。コンテスト終了後に非再帰を通した。以下ははじめに出した再帰のもの。混乱してしまい汚いコードになった。globalの使い方忘れてdictionaryの値をglobal変数として使ったり。

import copy
import sys
sys.setrecursionlimit(100000)
input = sys.stdin.readline

N,U,V = map(int, input().split())
d = [set() for _ in range(N+1)]
dps = [0 for _ in range(N+1)]
dpss = [0 for _ in range(N+1)]
vs = set()
r = {}

def it(v, p, depth, ps):
    for u in d[v]:
        if u == p:
            continue
        dps[u] = depth
        if u == V:
            r["p"] = copy.deepcopy(ps)
            continue
        ps.add(u)
        it(u, v, depth+1,ps)
        ps.remove(u)
        vs.add(u)

def it2(v, p, depth):
    for u in d[v]:
        if u == p:
            continue
        if u not in vs:
            continue
        if dps[v] >= depth-1 and u not in r["p"]:
            continue
        dpss[u] = depth
        it2(u, v, depth+1)

for _ in range(N-1):
    a, b = [int(x) for x in input().split()]
    d[a].add(b)
    d[b].add(a)

vs.add(U)
it(U, -1, 1, set([U]))
it2(V, -1, 1)
print(max(dpss)-1)

コンテスト後に非再帰にしたのは以下。

import sys
input = sys.stdin.readline
N,U,V = map(int, input().split())
d = [set() for _ in range(N+1)]
dps = [0 for _ in range(N+1)]
dpss = [0 for _ in range(N+1)]

def it(v):
    stack = [v]
    vs = set()
    path = []
    while True:
        if not len(stack):
            break
        v = stack.pop()
        if v > 0:
            path.append(v)
            vs.add(v)
            for u in d[v]:
                if u in vs:
                    continue
                if u == V:
                    r = path[:]
                dps[u] = dps[v] + 1
                stack.append(-v)
                stack.append(u)
        else:
            path.pop()
    return r

def it2(v, path):
    stack = [v]
    vs = set()
    while True:
        if not len(stack):
            break
        v = stack.pop()
        vs.add(v)
        for u in d[v]:
            if u in vs:
                continue
            if dps[v] >= dpss[v] and u not in path:
                continue
            dpss[u] = dpss[v] + 1
            stack.append(u)

for _ in range(N-1):
    a, b = [int(x) for x in input().split()]
    d[a].add(b)
    d[b].add(a)

path = set(it(U))
it2(V, path)
print(max(dpss)-1)

以下のpathの保持の部分が少し苦戦したので要復習。親ノード番号を*-1してstackにappendすることで深さ優先時、今の位置を保存できる。

    stack = [v]
    vs = set()
    path = []
    while True:
        if not len(stack):
            break
        v = stack.pop()
        if v > 0:
            path.append(v)
            vs.add(v)
            for u in d[v]:
                if u in vs:
                    continue
                if u == V:
                    r = path[:]
                    continue
                dps[u] = dps[v] + 1
                stack.append(-v)
                stack.append(u)
        else:
            path.pop()