#877589 (PyPy3) No.2337 Equidistant

• 問題: No.2337 Equidistant
• ユーザー: navel_tos
• 提出日時: 2023-06-02 23:42:49
• 言語: PyPy3 (7.3.15)
• 結果: WA
• 実行時間: -
• コード長: 6,099 bytes
• コンパイル時間: 499 ms
• コンパイル使用メモリ: 82,288 KB
• 実行使用メモリ: 540,128 KB
• 最終ジャッジ日時: 2024-12-29 00:04:19
• 合計ジャッジ時間: 48,374 ms
• ジャッジサーバーID （参考情報）: judge4 / judge2

ファイルパターン	結果
sample	AC * 1
other	AC * 7 WA * 20 TLE * 1

ソースコード

#yukicoder391D

'''
全方位木DPで気合いいれて解くしかないかなぁ。

まずこのXはSi,Tiの経路から分岐した点でないといけないよね？
だからこの経路上を探索しないといけない

二点の中点と、そこから分岐する頂点数がわかればよさそうだ。
LCAと全方位木DPで戦うか。きつい戦いになりそうだ。
一旦ライブラリを貼る。
'''
#最小共通祖先 ライブラリ
class LowestCommonAncestor:
    def __init__(self,n):
        self._n=n;n=0
        while 2**(n/10)<self._n:n+=1
        self._logn=int(n/10+2);self._depth=[0 for _ in [0]*self._n];self._distance=[0 for _ in [0]*self._n];self._ancestor=[[-1 for _ in [0]*self._n] for k in [0]*self._logn];self._edge=[[] for _ in [0]*self._n]
    def add_edge(self,u,v,w=1):self._edge[u].append((v,w));self._edge[v].append((u,w))
    def build(self,root=0):  #rootを指定し、その他の頂点に祖先情報を書き込む
        stack=[root]
        while stack:
            now=stack.pop()
            for nxt,w in self._edge[now]:
                if self._ancestor[0][nxt]!=now and self._ancestor[0][now]!=nxt:self._ancestor[0][nxt]=now;self._depth[nxt]=self._depth[now]+1;self._distance[nxt]=self._distance[now]+w;stack.append(nxt)
        for k in range(1,self._logn):
            for i in range(self._n):
                if self._ancestor[k-1][i]==-1:self._ancestor[k][i]=-1
                else:self._ancestor[k][i]=self._ancestor[k-1][self._ancestor[k-1][i]]
    def LCA(self,u,v):
        if self._depth[u]>self._depth[v]:u,v=v,u
        for k in range(self._logn-1,-1,-1):v=self._ancestor[k][v] if((self._depth[v]-self._depth[u])>>k)&1 else v
        if u==v:return u
        for k in range(self._logn-1,-1,-1):  #ギリギリ一致する直前まで祖先を辿る
            if self._ancestor[k][u]!=self._ancestor[k][v]:
                u,v=self._ancestor[k][u],self._ancestor[k][v]
        return self._ancestor[0][u]
    def distance(self,u,v):return self._distance[u]+self._distance[v]-2*self._distance[self.LCA(u,v)]


import sys; sys.setrecursionlimit(10**7); input=sys.stdin.readline
f=lambda:list(map(int,input().split()))

N,Q=f(); G=[[] for _ in range(N)]; LCA=LowestCommonAncestor(N)
for _ in range(N-1):
    a,b=f(); G[a-1].append(b-1); G[b-1].append(a-1); LCA.add_edge(a-1,b-1)

#面倒くさいので、親は次数1の頂点とする
for P in range(N):
    if len(G[P])==1: break
LCA.build(P)

'''
考察再開。
ここから全方位木DPを行えば、「どの方向の辺を辿るとどれだけの頂点があるか」は
即答できるようになる。問題は、どの辺はカウントが不要なのか考えないといけない。

結局DFSかなぁ。どこが争点となるかは事前に判別可能だから、それを使ってDFSかな。

まずは中点の判定を行おう。丁寧にやれば大丈夫そう。
'''
Task=[0]*Q; MidPoint=[-1]*Q; checkpoint=[[] for _ in range(N)]; MP=[set() for _ in range(N)]
for i in range(Q):
    a,b=f(); a-=1; b-=1; Task[i]=(a,b)
    x=LCA.LCA(a,b); da,db,dx=LCA._distance[a],LCA._distance[b],LCA._distance[x]
    da,db=da-dx,db-dx  #LCAからaの距離、LCAからbの距離を格納
    if (da+db)%2: continue
    mid=(da+db)//2
    if da>=mid:
        #aからmidだけ戻った頂点をMidPointに格納する
        midp=bin(mid)[2:]; midL=len(midp); now=a
    else: midp=bin(mid)[2:]; midL=len(midp); now=b
    for j in range(midL):
        midbit=midL-1-j
        if midp[j]=='1': now=LCA._ancestor[midbit][now]
    MidPoint[i]=now; checkpoint[Task[i][0]].append(i); checkpoint[Task[i][1]].append(i)
    if da==db: MP[now].add(i)
del LCA

            
'''
この前処理で、各クエリの中点は判定できた。

あり得るケースは
1. 経路上に2点があり、経路中に中点がある
2. 分岐上に2点がある
ケース。

DFSで「始点から何個頂点があるか」を保持しつつ、
帰りがけに「どの辺方向を辿ると、その先に何個頂点があるか」を返せばいいか。

違うな。
各中点クエリに対して、帰りがけに「この辺方向から来たやつは数えちゃだめよ」をやるのか。

1. 各頂点の「親方向を除く」頂点数
2. 各クエリの減算すべき頂点数
これらを別カウントしてDFS、きついけれどこれしかない。

つまり、中点に返ってきた方向の頂点数は数えてはいけない、というルール。
各クエリごとに番号を記録しておき、中点に戻るたびに印をdiscardする感じかな。

各頂点ごとに、中点として判定しなければならないクエリ番号を保持しておくか。
既にvisitedで、かつ戻ってきたときにのみ判定したい。(行きがけは判定しない)
できるか？できそう。
「戻ってきたらこれらのクエリ番号に対して、経路上の頂点数の減算を行え」みたいにやるか。

とりあえずTLEによる強制打ち切りは喰らわなかったので、方針はこれでよい。
後はデバッグ。

うわ、1ケースTLEしている。苦しい。
'''
#Vert[i]: 頂点iの、「親方向を除いた」頂点数(部分木の頂点数)
Vert=[0]*N; QueryV=[0]*Q; visited=[0]*N; duty=[set() for _ in range(N)]; Save=[0]*N

#DFS 部分木の頂点数を数えつつ、帰りがけにdutyを爆破する
def yukicoder391D(now,vertice):
    for i in MP[now]: QueryV[i]+=vertice
    visited[now]=1; vertice+=1; Save[now]=vertice; Vert[now]+=1
    for t in checkpoint[now]:
        if visited[MidPoint[t]]: duty[MidPoint[t]].add(t)
    for next in G[now]:
        if visited[next]: continue
        vertice=yukicoder391D(next,vertice)
        #帰りがけにtaskを処理
        v=vertice-Save[now]; Save[now]=vertice; Vert[now]+=v   #この経路上の頂点数
        for taskno in duty[now]: QueryV[taskno]-=v
        duty[now]=set()
    return vertice

yukicoder391D(P,0)
for i in range(Q):
    if MidPoint[i]==-1: print(0); continue
    print(Vert[MidPoint[i]]+QueryV[i])