結果

問題 No.778 クリスマスツリー
ユーザー 👑 Kazun
提出日時 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
権限があれば一括ダウンロードができます

ソースコード

diff #
プレゼンテーションモードにする

class Tree:
def __init__(self,N,index=0):
"""N(index, index+1, ..., N-1+index).
"""
self.N=N
self.index=index
self.parent=[-1]*(N+index)
self.__mutable=True
def vertex_exist(self,x):
return self.index<=x<self.index+self.N
def __after_seal_check(self,*vertexes):
if self.__mutable:
return False
for v in vertexes:
if not self.vertex_exist(v):
return False
return True
def is_mutable(self):
return self.__mutable
#
def root_set(self,root):
"""x.
"""
assert self.vertex_exist(root)
assert self.__mutable
self.root=root
def parent_set(self,x,y):
"""xy.
"""
assert self.vertex_exist(x)
assert self.vertex_exist(y)
assert self.__mutable
self.parent[x]=y
def child_set(self,x,y):
"""xy.
"""
assert self.vertex_exist(x)
assert self.vertex_exist(y)
assert self.__mutable
self.parent[y]=x
def seal(self):
""".
"""
assert self.__mutable
assert hasattr(self,"root")
a=self.index
b=self.index+self.N
C=[[] for _ in range(b)]
p=self.parent
ve=self.vertex_exist
for i in range(a,b):
if i!=self.root:
assert ve(p[i])
C[p[i]].append(i)
self.__mutable=False
self.children=C
#.
def depth_search(self,Mode=True):
""".
"""
assert self.__after_seal_check()
if hasattr(self,"depth"):
return self.depth
from collections import deque
C=self.children
D=[-1]*(self.index+self.N)
E=[[] for _ in range(self.N)]
Q=deque([self.root])
D[self.root]=0
E[0]=[self.root]
while Q:
x=Q.popleft()
d=D[x]
for y in C[x]:
D[y]=d+1
E[d+1].append(y)
Q.append(y)
self.depth=D
self.tower=E
if Mode:
return D
def 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"):
return
if 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.parent
rg=range(self.index,self.index+self.N)
for x in rg:
if x!=self.root:
Y[x]=p[x]
else:
Y[x]=self.root
for k in range(1,b):
Y=X[k-1]
Z=X[k]
for x in rg:
Z[x]=Y[Y[x]]
self.upper_list=X
def upper(self,x,k,over=True):
"""xk.
over:(x)<dTrue, False.
"""
assert self.__after_seal_check(x)
assert 0<=k
if not hasattr(self,"upper_list"):
self.__upper_list()
if self.vertex_depth(x)<k:
if over:
return self.root
else:
raise ValueError
i=0
while k:
if k&1:
x=self.upper_list[i][x]
k>>=1
i+=1
return x
def 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,x
dd=-dd
y=self.upper(y,dd)
if x==self.root:
return x
if x==y:
return x
d=self.vertex_depth(x)
b=d.bit_length()
X=self.upper_list
for k in range(b-1,-1,-1):
px=X[k][x];py=X[k][y]
if px!=py:
x=px;y=py
return self.upper(x,1)
def __degree_count(self):
assert self.__after_seal_check()
if hasattr(self,"deg"):
return
self.deg=[0]*(self.index+self.N)
for v in range(self.index,self.index+self.N):
d=len(self.children[v])+1
if d!=self.root:
d-=1
self.deg[v]=d
return
def 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 deque
def bfs(start):
X=[-1]*(self.index+self.N)
Q=deque([start])
X[start]=0
pa=self.parent
ch=self.children
while Q:
x=Q.popleft()
if X[pa[x]]==-1:
Q.append(pa[x])
X[pa[x]]=X[x]+1
for y in ch[x]:
if X[y]==-1:
Q.append(y)
X[y]=X[x]+1
y=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,d
def path(self,u,v):
"""u,v.
"""
assert self.__after_seal_check(u,v)
w=self.lowest_common_ancestor(u,v)
pa=self.parent
X=[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):
"""2u,v () ?
"""
assert self.__after_seal_check(u,v)
if u==self.root or v==self.root:
return False
return self.parent[u]==self.parent[v]
def is_ancestor(self,u,v):
"""uv?
"""
assert self.__after_seal_check(u,v)
dd=self.vertex_depth(v)-self.vertex_depth(u)
if dd<0:
return False
v=self.upper(v,dd)
return u==v
def is_descendant(self,u,v):
"""uv?
"""
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):
"""2u,v.
"""
assert self.__after_seal_check(u,v)
dep=self.vertex_depth
return 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"):
return
if not hasattr(self,"tower"):
self.depth_search(False)
self.des_count=[1]*(self.index+self.N)
pa=self.parent
for T in self.tower[:0:-1]:
for x in T:
self.des_count[pa[x]]+=self.des_count[x]
return
def 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 deque
Q=deque([self.root])
T=[-1]*(self.N+self.index)
p=1
while Q:
x=Q.popleft()
T[x]=p
p+=1
C=self.children[x]
for y in C:
Q.append(y)
self.preorder_number=T
return T[v]
def dfs_yielder(self):
"""DFSyield.
(v,1): v
(v,0): v
"""
assert self.__after_seal_check()
#
yield (self.root,1)
v=self.root
ch=self.children
pa=self.parent
R=[-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:
break
else:
v=pa[v]
else: #
w=v
v=ch[v][S[v]]
S[w]+=1
yield (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 v
def 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 v
def tree_dp(self,calc,unit,f,g,Mode=False):
"""DP.
[input]
calc:2 M x M -> M
unit:M
f,g: M -> M
Mode: False->, True->
[]
v x,y,z,...,
dp[v]=g(f(x)*f(y)*f(z)*...)
.
"""
data=[unit]*(self.index+self.N)
pa=self.parent
for x in self.bottom_up():
if x==self.root:
break
data[x]=g(data[x])
data[pa[x]]=calc(data[pa[x]],f(data[x]))
if Mode:
return data
else:
return data[self.root]
#==================================================
class Binary_Indexed_Tree():
def __init__(self,L,calc,unit,inv,index=1):
"""calcNBinary Indexed Tree
calc:(2,)
unit:calc(xe=ex=xe)
inv:calc(1)
"""
self.calc=calc
self.unit=unit
self.inv=inv
self.index=index
N=len(L)
d=max(1,(N-1).bit_length())
k=2**d
X=[None]+[unit]*k
self.num=k
self.depth=d
if L:
for i in range(len(L)):
p=i+1
while p<=k:
X[p]=self.calc(X[p],L[i])
p+=p&(-p)
self.data=X
def index_number(self,k,index=1):
"""k.
k:
index:
"""
return self.sum(k,k,index)
def add(self,k,x,index=1,right=False):
"""kx,.
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):
"""kx,.
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):
"""FromTo.
※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.unit
x=To
while x>0:
S=self.calc(self.data[x],S)
x-=x&(-x)
return S
def 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,neg
N=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=0
for 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)
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0