#877726 (PyPy3) No.2337 Equidistant

提出ソース
結果

問題	No.2337 Equidistant
コンテスト
ユーザー	navel_tos
提出日時	2023-06-03 01:12:57
言語	PyPy3 (7.3.15)
結果	MLE
実行時間	-
コード長	7,013 bytes
コンパイル時間	295 ms
コンパイル使用メモリ	82,344 KB
実行使用メモリ	1,193,712 KB
最終ジャッジ日時	2024-12-29 02:01:51
合計ジャッジ時間	76,536 ms
ジャッジサーバーID （参考情報）	judge4 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
sample	MLE * 1
other	AC * 22 TLE * 5 MLE * 1
権限があれば一括ダウンロードができます
ソースコード

raw source code
#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
from collections import deque as dq
f=lambda:list(map(int,input().split()))

#全方位木DPをしたい G[i]: (移動方向、逆辺)
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,len(G[b-1]))); G[b-1].append((a-1,len(G[a-1])-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=[dq() 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)
del LCA,Task
            
'''
この前処理で、各クエリの中点は判定できた。

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

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

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

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

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

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

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

うわ、1ケースTLEしている。苦しい。

わかった。ちゃんと全方位木DPをしないとだめだったんだ･･･

■AfterContest
くやしい　やるか

一旦頂点数だけを見る全方位木DPをやろう
'''
#Vert[i][j]: 頂点iの、辺j方向の頂点数
Vert=[[-1]*len(G[i]) for i in range(N)]; visited=[1]+[0]*N

#頂点数の格納のみに集中
def yukicoder391D(now,back,vertice):
    vertice+=1; save=vertice; visited[visited[0]]=now; visited[0]+=1
    for index,(next,_) in enumerate(G[now]):
        if next==back: continue
        vertice=yukicoder391D(next,now,vertice)
        Vert[now][index]=vertice-save; save=vertice
    return vertice
yukicoder391D(P,-1,0)

#トポロジカルソート順に全方位木DP
root=[-1]*N
for now in visited[1:]:
    sV=sum(Vert[now])
    for index,(next,rev) in enumerate(G[now]):
        if root[next]==-1: root[next]=(next,rev)
        Vert[next][rev]=sV-Vert[now][index]+1
        
#QueryV[i]: クエリiの、勘定してはならない(頂点,辺番号)
QueryV=[[] for _ in range(Q)]
        
#数えてはいけない辺を探すDFS 帰りがけにだけ処理すればよい
visited=[0]*N; duty=[dq() for _ in range(N)]
def yukicoder391D_2(now,back):
    visited[now]=1
    while checkpoint[now]:
        t=checkpoint[now].pop()
        if visited[MidPoint[t]]: duty[MidPoint[t]].append(t)
    for index,(next,_) in enumerate(G[now]):
        if next==back: continue
        yukicoder391D_2(next,now)
        while duty[now]: t=duty[now].pop(); QueryV[t].append((now,index))
yukicoder391D_2(P,-1); del duty
        
#クエリに回答
for i in range(Q):
    if MidPoint[i]<0: print(0); continue
    if len(QueryV[i])<2: QueryV[i].append(root[MidPoint[i]])
    ans=N
    for v,e in QueryV[i]: ans-=Vert[v][e]
    print(ans)