結果
問題 | No.778 クリスマスツリー |
ユーザー |
👑 ![]() |
提出日時 | 2021-05-14 03:37:57 |
言語 | PyPy3 (7.3.15) |
結果 |
AC
|
実行時間 | 616 ms / 2,000 ms |
コード長 | 13,957 bytes |
コンパイル時間 | 197 ms |
コンパイル使用メモリ | 82,700 KB |
実行使用メモリ | 118,488 KB |
最終ジャッジ日時 | 2024-10-01 07:42:19 |
合計ジャッジ時間 | 6,042 ms |
ジャッジサーバーID (参考情報) |
judge5 / judge1 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 3 |
other | AC * 12 |
ソースコード
class Tree: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):return self.index<=x<self.index+self.Ndef __after_seal_check(self,*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"):return self.depthfrom 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の深さ)<dのときに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 y,z,ddef 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):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):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(self,calc,unit,f,g,Mode=False):"""葉から木DPを行う.[input]calc:モノイドを成す2項演算 M x M -> Munit:Mの単位元f,g: M -> MMode: False->根の値のみ, True->全ての値[補足]頂点 v の子が x,y,z,...のとき, 更新式はdp[v]=g(f(x)*f(y)*f(z)*...)になる."""data=[unit]*(self.index+self.N)pa=self.parentfor x in self.bottom_up():if x==self.root:breakdata[x]=g(data[x])data[pa[x]]=calc(data[pa[x]],f(data[x]))if Mode:return dataelse:return data[self.root]#==================================================class Binary_Indexed_Tree():def __init__(self,L,calc,unit,inv,index=1):"""calcを演算とするN項のBinary Indexed Treeを作成calc:演算(2変数関数,群)unit:群calcの単位元(xe=ex=xを満たすe)inv:群calcの逆元(1変数関数)"""self.calc=calcself.unit=unitself.inv=invself.index=indexN=len(L)d=max(1,(N-1).bit_length())k=2**dX=[None]+[unit]*kself.num=kself.depth=dif L:for i in range(len(L)):p=i+1while p<=k:X[p]=self.calc(X[p],L[i])p+=p&(-p)self.data=Xdef index_number(self,k,index=1):"""第k要素の値を出力する.k:数列の要素index:先頭の要素の番号"""return self.sum(k,k,index)def add(self,k,x,index=1,right=False):"""第k要素にxを左から加え,更新を行う.k:数列の要素x:更新後の値index:先頭の要素の番号right:「左から」が「右から」になる"""p=k+(1-index)while p<=self.num:if right==False:#左からself.data[p]=self.calc(x,self.data[p])else:#右からself.data[p]=self.calc(self.data[p],x)p+=p&(-p)def update(self,k,x,index=1,right=False):"""第k要素をxに変え,更新を行う.k:数列の要素x:更新後の値"""a=self.index_number(k,index)if right==False:#左からy=self.calc(x,self.inv(a))else:#右からy=self.calc(self.inv(a),x)self.add(k,y,index,right)def sum(self,From,To,index=1):"""第From要素から第To要素までの総和を求める.※From!=1を使うならば,群でなくてはならない.From:始まりTo:終わりindex:先頭の要素の番号"""alpha=max(1,From+(1-index))beta=min(self.num,To+(1-index))if alpha==1:return self.__section(beta)else:return self.calc(self.inv(self.__section(alpha-1)),self.__section(beta))def __section(self,To):S=self.unitx=Towhile x>0:S=self.calc(self.data[x],S)x-=x&(-x)return Sdef all_sum(self):return self.data[-1]def __getitem__(self,index):if isinstance(index,int):return self.index_number(index,self.index)else:return [self.index_number(t,self.index) for t in index]def __setitem__(self,index,val):self.update(index,val,self.index)#=================================================from operator import add,negN=int(input())A=[-1]+list(map(int,input().split()))T=Tree(N)T.root_set(0)for i in range(1,N):T.parent_set(i,A[i])T.seal()B=Binary_Indexed_Tree([0]*N,add,0,neg,0)X=0for v,m in T.dfs_yielder():if m==1:X+=B.sum(0,v,0)B.add(v,1,0)else:B.add(v,-1,0)print(X)