結果
問題 | No.1333 Squared Sum |
ユーザー |
![]() |
提出日時 | 2023-02-18 02:21:09 |
言語 | PyPy3 (7.3.15) |
結果 |
AC
|
実行時間 | 1,984 ms / 2,000 ms |
コード長 | 17,621 bytes |
コンパイル時間 | 176 ms |
コンパイル使用メモリ | 82,848 KB |
実行使用メモリ | 298,920 KB |
最終ジャッジ日時 | 2024-07-19 16:47:06 |
合計ジャッジ時間 | 49,134 ms |
ジャッジサーバーID (参考情報) |
judge5 / judge3 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
other | AC * 44 |
ソースコード
#https://atcoder.jp/contests/abc222/submissions/26476209class Tree:__slots__=("N", "index", "parent", "__mutable","root", "children", "depth", "tower", "upper_list", "deg", "des_count", "preorder_number","euler", "in_time", "out_time")def __init__(self,N,index=0):""" N 頂点 (index, index+1, ..., N-1+index) の根付き木を生成する. """self.N=Nself.index=indexself.parent=[-1]*(N+index)self.__mutable=Truedef vertex_exist(self,x):""" 頂点 x が存在するかどうかを判定する. """return self.index<=x<self.index+self.Ndef __after_seal_check(self,*vertexes):""" 木が確定していて, vertexes の頂点が存在するかどうかをチェックする. """if self.__mutable:return Falsefor v in vertexes:if not self.vertex_exist(v):return Falsereturn Truedef is_mutable(self):""" 木が確定して [いない] かどうかを返す. """return self.__mutable#設定パートdef root_set(self,root):""" 頂点 x を根に設定する."""assert self.vertex_exist(root)assert self.__mutableself.root=rootdef parent_set(self,x,y):""" 頂点 x の親を y に設定する."""assert self.vertex_exist(x)assert self.vertex_exist(y)assert self.__mutableself.parent[x]=ydef child_set(self,x,y):""" 頂点 x の子の一つに y を設定する."""assert self.vertex_exist(x)assert self.vertex_exist(y)assert self.__mutableself.parent[y]=xdef seal(self):""" 木の情報を確定させる."""assert self.__mutableassert hasattr(self,"root")a=self.indexb=self.index+self.NC=[[] for _ in range(b)]p=self.parentve=self.vertex_existfor i in range(a,b):if i!=self.root:assert ve(p[i])C[p[i]].append(i)self.__mutable=Falseself.children=C#データを求める.def depth_search(self,Mode=True):""" 木の深さを求める. """assert self.__after_seal_check()if hasattr(self,"depth"):if Mode:return self.depthelse:returnfrom collections import dequeC=self.childrenD=[-1]*(self.index+self.N)E=[[] for _ in range(self.N)]Q=deque([self.root])D[self.root]=0E[0]=[self.root]while Q:x=Q.popleft()d=D[x]for y in C[x]:D[y]=d+1E[d+1].append(y)Q.append(y)self.depth=Dself.tower=Eif Mode:return Ddef vertex_depth(self,x):""" 頂点 x の深さを求める."""assert self.__after_seal_check(x)if not hasattr(self,"depth"):self.depth_search(Mode=False)return self.depth[x]def __upper_list(self):assert self.__after_seal_check()if hasattr(self,"upper_list"):returnif not hasattr(self,"depth"):self.depth_search(False)b=max(self.depth).bit_length()X=[[-1]*(self.index+self.N) for _ in range(b)]Y=X[0]p=self.parentrg=range(self.index,self.index+self.N)for x in rg:if x!=self.root:Y[x]=p[x]else:Y[x]=self.rootfor k in range(1,b):Y=X[k-1]Z=X[k]for x in rg:Z[x]=Y[Y[x]]self.upper_list=Xdef upper(self,x,k,over=True):""" 頂点 x から見て k 個親の頂点を求める.over: (頂点 x の深さ)<k のときに True ならば根を返し, False ならばエラーを吐く."""assert self.__after_seal_check(x)assert 0<=kif not hasattr(self,"upper_list"):self.__upper_list()if self.vertex_depth(x)<k:if over:return self.rootelse:raise ValueErrori=0while k:if k&1:x=self.upper_list[i][x]k>>=1i+=1return xdef lowest_common_ancestor(self,x,y):""" 頂点 x, y の最小共通先祖 (x,yに共通する先祖で最も深いもの) を求める. """assert self.__after_seal_check(x,y)dd=self.vertex_depth(y)-self.vertex_depth(x)if dd<0:x,y=y,xdd=-ddy=self.upper(y,dd)if x==self.root:return xif x==y:return xd=self.vertex_depth(x)b=d.bit_length()X=self.upper_listfor k in range(b-1,-1,-1):px=X[k][x];py=X[k][y]if px!=py:x=px;y=pyreturn self.upper(x,1)def __degree_count(self):assert self.__after_seal_check()if hasattr(self,"deg"):returnself.deg=[0]*(self.index+self.N)for v in range(self.index,self.index+self.N):d=len(self.children[v])+1if d==self.root:d-=1self.deg[v]=dreturndef degree(self,v):""" 頂点 v の次数を求める. """assert self.__after_seal_check(v)if not hasattr(self,"deg"):self.__degree_count()return self.deg[v]def diameter(self):""" 木の直径を求める."""assert self.__after_seal_check()from collections import dequedef bfs(start):X=[-1]*(self.index+self.N)Q=deque([start])X[start]=0pa=self.parentch=self.childrenwhile Q:x=Q.popleft()if X[pa[x]]==-1:Q.append(pa[x])X[pa[x]]=X[x]+1for y in ch[x]:if X[y]==-1:Q.append(y)X[y]=X[x]+1y=max(range(self.index,self.index+self.N),key=lambda x:X[x])return y,X[y]y,_=bfs(self.root)z,d=bfs(y)return d,(y,z)def path(self,u,v):""" 頂点 u, v 間のパスを求める. """assert self.__after_seal_check(u,v)w=self.lowest_common_ancestor(u,v)pa=self.parentX=[u]while u!=w:u=pa[u]X.append(u)Y=[v]while v!=w:v=pa[v]Y.append(v)return X+Y[-2::-1]def is_brother(self,u,v):""" 2つの頂点 u, v は兄弟 (親が同じ) か? """assert self.__after_seal_check(u,v)if u==self.root or v==self.root:return Falsereturn self.parent[u]==self.parent[v]def is_ancestor(self,u,v):""" 頂点 u は頂点 v の先祖か? """assert self.__after_seal_check(u,v)dd=self.vertex_depth(v)-self.vertex_depth(u)if dd<0:return Falsev=self.upper(v,dd)return u==vdef is_descendant(self,u,v):""" 頂点 u は頂点 v の子孫か? """assert self.__after_seal_check(u,v)return self.is_ancestor(v,u)def is_leaf(self,v):""" 頂点 v は葉? """return not bool(self.children[v])def distance(self,u,v):""" 2頂点 u, v 間の距離を求める. """assert self.__after_seal_check(u,v)dep=self.vertex_depthreturn dep(u)+dep(v)-2*dep(self.lowest_common_ancestor(u,v))def __descendant_count(self):assert self.__after_seal_check()if hasattr(self,"des_count"):returnif not hasattr(self,"tower"):self.depth_search(False)self.des_count=[1]*(self.index+self.N)pa=self.parentfor T in self.tower[:0:-1]:for x in T:self.des_count[pa[x]]+=self.des_count[x]returndef descendant_count(self, v):""" 頂点 v の子孫の数を求める. """assert self.__after_seal_check(v)self.__descendant_count()return self.des_count[v]def subtree_size(self, v):""" 頂点 v を根とした部分根付き木のサイズを求める. """return self.descendant_count(v)def preorder(self,v):""" 頂点 v の行きがけ順を求める. """assert self.__after_seal_check(v)if hasattr(self,"preorder_number"):self.preorder_number[v]from collections import dequeQ=deque([self.root])T=[-1]*(self.N+self.index)p=1while Q:x=Q.popleft()T[x]=pp+=1C=self.children[x]for y in C:Q.append(y)self.preorder_number=Treturn T[v]def dfs_yielder(self):""" DFS における頂点の出入りを yield する.(v,1): 頂点 v に入る(v,0): 頂点 v を出る"""assert self.__after_seal_check()#最初yield (self.root,1)v=self.rootch=self.childrenpa=self.parentR=[-1]*self.index+[len(ch[x]) for x in range(self.index,self.index+self.N)]S=[0]*(self.index+self.N)while True:if R[v]==S[v]: #もし,進めないならばyield (v,0) #頂点vを出るif v==self.root:breakelse:v=pa[v]else: #進めるw=vv=ch[v][S[v]]S[w]+=1yield (v,1)def top_down(self):""" 木の頂点から yield する. """assert self.__after_seal_check()if not hasattr(self,"tower"):self.depth_search(False)for E in self.tower:for v in E:yield vdef bottom_up(self):""" 木の根から yield する. """assert self.__after_seal_check()if not hasattr(self,"tower"):self.depth_search(False)for E in self.tower[::-1]:for v in E:yield vdef tree_dp_from_leaf(self,merge,unit,f,g,Mode=False):""" 葉から木 DP 行う.[input]merge: 可換モノイドを成す2項演算 M x M -> Munit: Mの単位元f: X x V x V → M: f(x,v,w): v が親, w が子g: M x V → X: g(x,v)Mode: False → 根の値のみ, True → 全ての値[補足]頂点 v の子が x,y,z,...のとき, 更新式は * を merge としてdp[v]=g(f(dp[x],v,x)*f(dp[y],v,y)*f(dp[z],v,z)*..., v)になる."""assert self.__after_seal_check()data=[unit]*(self.index+self.N)ch=self.childrenfor x in self.bottom_up():for y in ch[x]:data[x]=merge(data[x],f(data[y],x,y))data[x]=g(data[x],x)if Mode:return dataelse:return data[self.root]def tree_dp_from_root(self,f,alpha):""" 根から木 DP を行う.[input]alpha: 初期値f: X x V x V -> X: f(x,v,w): v が親, w が子[補足]頂点 v の親が x のとき, 更新式はdp[v]=f(dp[x],x,v) (x!=root), alpha (x==root)になる."""assert self.__after_seal_check()data=[0]*(self.index+self.N)ch=self.childrendata[self.root]=alphafor x in self.top_down():for y in ch[x]:data[y]=f(data[x],x,y)return datadef rerooting(self,merge,unit,f,g):""" 全方位木 DP を行う.[input]calc:可換モノイドを成す2項演算 M x M -> Munit:Mの単位元f: X x V x V -> M: f(x,v,w): v が親, w が子g: M x V -> X: g(x,v)※ tree_dp_from_leaf と同じ形式[補足]頂点 v の子が x,y,z,...のとき, 更新式はdp[v]=g(f(dp[x],v,x)*f(dp[y],v,y)*f(dp[z],v,z)*..., v)になる."""assert self.__after_seal_check()upper=[unit]*(self.index+self.N)lower=[unit]*(self.index+self.N)ch=self.childrenpa=self.parent#DFSパートlower=self.tree_dp_from_leaf(merge,unit,f,g,True)#BFSパートfor v in self.top_down():cc=ch[v]#累積マージdeg=len(cc)Left=[unit]; x=unitfor c in cc:x=merge(x,f(lower[c],v,c))Left.append(x)Right=[unit]; y=unitfor c in cc[::-1]:y=merge(y,f(lower[c],v,c))Right.append(y)Right=Right[::-1]for i in range(deg):c=cc[i]a=merge(Left[i],Right[i+1])if v!=self.root:b=merge(a,f(upper[v],v,pa[v]))else:b=aupper[c]=g(b,v)A=[unit]*(self.index+self.N)for v in range(self.index,self.index+self.N):if v!=self.root:a=f(upper[v],v,pa[v])else:a=unitfor c in ch[v]:a=merge(a,f(lower[c],v,c))A[v]=g(a,v)return Adef euler_tour(self):""" オイラーツアーに関する計算を行う. """assert self.__after_seal_check()if hasattr(self,"euler"):return#最初X=[]; X_append=X.append #X: Euler Tour のリストv=self.rootch=self.childrenpa=self.parentR=[-1]*self.index+[len(ch[x]) for x in range(self.index,self.index+self.N)]S=[0]*(self.index+self.N)while True:X_append(v)if R[v]==S[v]: #もし,進めないならばif v==self.root:breakelse:v=pa[v]else: #進めるw=vv=ch[v][S[v]]S[w]+=1self.euler=Xself.in_time=[-1]*(self.index+self.N)self.out_time=[-1]*(self.index+self.N)for i in range(len(X)):v=X[i]if self.in_time[v]==-1:self.in_time[v]=self.out_time[v]=ielse:self.out_time[v]=idef centroid(self, all=False):""" 木の重心を求めるall: False → 重心のうちの1頂点. True → 全ての重心."""assert self.__after_seal_check()M=self.N//2if not hasattr(self,"des_count"):self.__descendant_count()G=[]; ch=self.children; des=self.des_countfor v in range(self.index, self.index+self.N):if self.N-des[v]>M:breakflag=1for x in ch[v]:if des[x]>M:flag=0breakif flag:if all:G.append(v)else:return vreturn G#=================================================def Making_Tree(N,E,root,index=0):"""木を作る.N:頂点数E: 辺のリストroot: 根"""from collections import dequeF=[[] for _ in range(index+N)]for u,v in E:assert index<=u<index+Nassert index<=v<index+Nassert u!=vF[u].append(v)F[v].append(u)X=[-1]*(index+N)X[root]=rootC=[[] for _ in range(index+N)]Q=deque([root])while Q:x=Q.popleft()for y in F[x]:if X[y]==-1:X[y]=xQ.append(y)C[x].append(y)T=Tree(N,index)T.root_set(root)T.parent=XT.children=CT.seal()return T"""EDC-v 例N,m=map(int, input().split())E=[]for _ in range(N-1):a,b=map(int,input().split())E.append((a,b))T=Making_Tree(N,E,1,1)for a in T.rerooting(lambda x,y:(x*y)%m,1,lambda x,y,z:x+1,lambda x,y:x)[1:]:print(a)""""""ABC220-F 例N=int(input())E=[]for _ in range(N-1):a,b=map(int,input().split())E.append((a,b))T=Making_Tree(N,E,1,1)for a in T.rerooting(lambda x,y:(x[0]+y[0],x[1]+y[1]),(0,0),lambda x,y,z:(x[0]+x[1],x[1]),lambda x,y:(x[0],x[1]+1))[1:]:print(a[0])"""def main():N=int(input())mod=10**9+7E=[]dists=dict()for _ in range(N-1):a,b,c=map(int,input().split())E.append((a,b))dists[(a<<18)|b]=cdists[(b<<18)|a]=cT=Making_Tree(N,E,1,1)def f(x,y,z):d = dists[(y<<18)|z]return (x[0]+((2*x[1]+x[2]*d)*d),x[1]+x[2]*d,x[2])ans=0for a in T.rerooting(lambda x,y:(x[0]+y[0],x[1]+y[1],x[2]+y[2]),(0,0,0),f,lambda x,y:(x[0]%mod,x[1]%mod,x[2]+1))[1:]:ans+=a[0]print(ans*pow(2,mod-2,mod)%mod)if __name__ == '__main__':main()