結果

問題 No.2654 [Cherry 6th Tune] Re: start! (Black Sheep)
ユーザー 👑 Kazun
提出日時 2023-08-01 01:12:45
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 3,521 ms / 7,000 ms
コード長 27,663 bytes
コンパイル時間 309 ms
コンパイル使用メモリ 82,176 KB
実行使用メモリ 202,572 KB
最終ジャッジ日時 2024-11-06 14:45:46
合計ジャッジ時間 72,481 ms
ジャッジサーバーID
(参考情報)
judge5 / judge2
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
other AC * 40
権限があれば一括ダウンロードができます

ソースコード

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

class Tree:
__slots__=("N", "index", "parent", "__mutable",
"root", "children", "depth", "tower", "upper_list", "des_count", "preorder_number",
"euler_vertex", "euler_edge", "in_time", "out_time", "lca_dst",
"hld_hedge")
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):
""" x . """
return self.index<=x<self.index+self.N
def __after_seal_check(self,*vertexes):
""" , 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):
""" root ."""
assert self.vertex_exist(root)
assert self.__mutable
self.root=root
def parent_set(self,x,y):
""" x y ."""
assert self.vertex_exist(x)
assert self.vertex_exist(y)
assert self.__mutable
self.parent[x]=y
def child_set(self, x, y):
""" x y ( x )."""
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):
""" .
mode=True , ."""
assert self.__after_seal_check()
if hasattr(self, "depth"):
if mode:
return self.depth
else:
return
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
for x in range(self.index,self.index+self.N):
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 range(self.index,self.index+self.N):
Z[x]=Y[Y[x]]
self.upper_list=X
def upper(self, x, k, over=True):
""" x k .
over: ( x )<k True , 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_greedy(self, x, y):
""" x, y (x,y) "" ."""
assert self.__after_seal_check(x,y)
dx=self.vertex_depth(x); dy=self.vertex_depth(y)
if dx<dy:
dx,dy=dy,dx
x,y=y,x
pa=self.parent
while dx>dy:
x=pa[x]
dx-=1
while x!=y:
x=pa[x]
y=pa[y]
return x
def __lca_prepare(self):
assert self.__after_seal_check()
N=self.N
bit=max(1, ((2*N-1)-1).bit_length())
D=[[0]*(2*N-1) for _ in range(bit)]
self.euler_tour_vertex()
tour=self.euler_vertex
D[0]=tour.copy()
dep=self.depth_search(True)
for i in range(1, bit):
shift=1<<i
tab=D[i]
for j in range(0, 2*N-1, 2*shift):
t=min(j+shift, 2*N-1)
tab[t-1]=tour[t-1]
for k in range(t-2, j-1, -1):
if dep[tour[k]]<dep[tab[k+1]]:
tab[k]=tour[k]
else:
tab[k]=tab[k+1]
if 2*N-1<=t:
break
tab[t]=tour[t]
r=min(t+shift, 2*N-1)
for k in range(t+1, r):
if dep[tab[k-1]]<dep[tour[k]]:
tab[k]=tab[k-1]
else:
tab[k]=tour[k]
self.lca_dst=D
return
def lowest_common_ancestor(self, x, y):
""" x, y (x,y) "" . """
assert self.__after_seal_check(x,y)
if not hasattr(self, "lca_dst"):
self.__lca_prepare()
a=self.in_time[x]; b=self.in_time[y]
if a>b:
x,y=y,x
a,b=b,a
if a==b:
return self.lca_dst[0][a]
p=(a^b).bit_length()-1
tab=self.lca_dst[p]
u=tab[a]; v=tab[b]
return u if self.vertex_depth(u)<self.vertex_depth(v) else v
def degree(self,v):
""" v . """
assert self.__after_seal_check(v)
if v==self.root:
return len(self.children[v])
else:
return len(self.children[v])+1
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 d,(y,z)
def path(self, u, v, faster=False):
""" u, v . """
assert self.__after_seal_check(u,v)
if faster:
w=self.lowest_common_ancestor(u,v)
else:
w=self.lowest_common_ancestor_greedy(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_parent(self, u, v):
""" u v ? """
assert self.__after_seal_check(u,v)
return v!=self.root and u==self.parent[v]
def is_children(self, u, v):
""" u v ? """
assert self.__after_seal_check(u,v)
return self.is_parent(v,u)
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 False
return 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 False
v=self.upper(v,dd)
return u==v
def is_descendant(self,u,v):
""" u v ? """
assert self.__after_seal_check(u,v)
return self.is_ancestor(v,u)
def direction(self, u, v):
""" u v (u!=v) u """
assert self.__after_seal_check(u,v)
assert u!=v
if self.is_ancestor(u,v):
du=self.vertex_depth(u)
dv=self.vertex_depth(v)
return self.upper(v,dv-(du+1))
else:
return self.parent[u]
def jump(self, u, v, k, default=-1):
""" u v k (0-indexed) ( k default)
u: int
v: int
k: int
default=-1: int
"""
assert self.__after_seal_check(u,v)
if k==0:
return u
# lca .
x=u; y=v
dx=self.vertex_depth(x); dy=self.vertex_depth(y)
if dx>dy:
x,y=y,x
dx,dy=dy,dx
y=self.upper(y, dy-dx)
if x==self.root or x==y:
w=x
else:
bit=dx.bit_length()
X=self.upper_list
for t in range(bit-1,-1,-1):
px=X[t][x]; py=X[t][y]
if px!=py:
x=px; y=py
w=self.parent[x]
dist_uw=self.vertex_depth(u)-self.vertex_depth(w)
dist_wv=self.vertex_depth(v)-self.vertex_depth(w)
if dist_uw+dist_wv<k:
return default
elif k<=dist_uw:
return self.upper(u, k)
else:
return self.upper(v, (dist_uw+dist_wv)-k)
def is_leaf(self,v):
""" v ? """
return not bool(self.children[v])
def distance(self, u, v, faster=True):
""" 2 u, v . """
assert self.__after_seal_check(u,v)
dep=self.vertex_depth
if faster:
return dep(u)+dep(v)-2*dep(self.lowest_common_ancestor(u,v))
else:
return dep(u)+dep(v)-2*dep(self.lowest_common_ancestor_greedy(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, order=None):
""" DFS yield .
() .
def dfs(v):
yield (v,1) # v
for w in self.children[v]:
dfs(w) # v .
yield (v,-1)
order (1): for w in self.children[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)
if order!=None:
for w in range(self.index, self.index+self.N):
ch[w].sort(key=order)
while True:
if R[v]==S[v]: #,
yield (v,-1) #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):
""" 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 v
def 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 v
def tree_dp_from_leaf(self,merge,unit,f,g,Mode=False):
""" DP .
[input]
merge: 2 M x M -> M
unit: 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,..., w , * merge
dp[v]=g(f(dp[x],v,x)*f(dp[y],v,y)*f(dp[z],v,z)*...*f(dp[w],v,w), v)
.
"""
assert self.__after_seal_check()
data=[unit]*(self.index+self.N)
ch=self.children
for 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 data
else:
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.children
data[self.root]=alpha
for x in self.top_down():
for y in ch[x]:
data[y]=f(data[x],x,y)
return data
def rerooting(self, merge, unit, f, g):
""" DP .
[input]
merge: 2 M x M -> M
unit: 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,..., w , * merge
dp[v]=g(f(dp[x],v,x)*f(dp[y],v,y)*f(dp[z],v,z)*...*f(dp[w],v,w), v)
.
"""
assert self.__after_seal_check()
upper=[unit]*(self.index+self.N)
lower=[unit]*(self.index+self.N)
ch=self.children
pa=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=unit
for c in cc:
x=merge(x, f(lower[c], v, c))
Left.append(x)
Right=[unit]; y=unit
for 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=a
upper[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=unit
for c in ch[v]:
a=merge(a, f(lower[c], v, c))
A[v]=g(a, v)
return A
def euler_tour_vertex(self, order=None):
""" (vertex) .
order: ()
"""
assert self.__after_seal_check()
if hasattr(self,"euler_vertex"):
return
#
X=[-1]*(2*self.N-1) #X: Euler Tour (vertex)
v=self.root
ch=self.children
if order!=None:
for i in range(self.index,self.index+self.N):
ch[i].sort(key=order)
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)
for t in range(2*self.N-1):
X[t]=v
if R[v]==S[v]:
v=pa[v]
else: #
w=v
v=ch[v][S[v]]
S[w]+=1
self.euler_vertex=X
self.in_time=[-1]*(self.index+self.N)
self.out_time=[-1]*(self.index+self.N)
for t in range(2*self.N-1):
v=X[t]
if self.in_time[v]==-1:
self.in_time[v]=self.out_time[v]=t
else:
self.out_time[v]=t
def euler_tour_edge(self):
""" (edge) .
(u, v, k): u v (k=+1 , k=-1 )
"""
assert self.__after_seal_check()
if hasattr(self,"euler_edge"):
return
if not hasattr(self, "euler_vertex"):
self.euler_tour_vertex()
self.euler_edge=[0]*(2*(self.N-1))
euler=self.euler_vertex
pa=self.parent
for t in range(2*(self.N-1)):
u=euler[t]; v=euler[t+1]
k=1 if u==pa[v] else -1
self.euler_edge[t]=(u,v,k)
def centroid(self, all=False):
"""
all: False → 1. True → .
"""
assert self.__after_seal_check()
M=self.N//2
if not hasattr(self,"des_count"):
self.__descendant_count()
G=[]; ch=self.children; des=self.des_count
for v in range(self.index, self.index+self.N):
if self.N-des[v]>M:
break
flag=1
for x in ch[v]:
if des[x]>M:
flag=0
break
if flag:
if all:
G.append(v)
else:
return v
return G
def generated_subtree(self,S):
""" S . """
assert self.__after_seal_check(*S)
if not hasattr(self, "in_time"):
self.euler_tour_vertex()
S=sorted(set(S),key=lambda i:self.in_time[i])
K=len(S)
T=set()
for i in range(K-1):
for a in self.path(S[i],S[i+1]):
T.add(a)
return sorted(T)
def generated_subtree_size(self,S):
""" S . """
assert self.__after_seal_check(*S)
if not hasattr(self, "in_time"):
self.euler_tour_vertex()
S=sorted(set(S),key=lambda i:self.in_time[i])
K=len(S)
X=0
for i in range(K-1):
X+=self.distance(S[i],S[i+1])
return (X+self.distance(S[-1],S[0]))//2
def Making_Tree_from_Adjacent_List(N, A, root, index=0):
""" ."""
from collections import deque
T=Tree(N, index)
T.root_set(root)
S=[False]*(N+index); S[root]=True
Q=deque([root])
while Q:
v=Q.popleft()
for w in A[v]:
if not S[w]:
S[w]=True
T.parent_set(w,v)
Q.append(w)
T.seal()
return T
#==================================================
class Binary_Indexed_Tree():
def __init__(self, L, op, zero, neg):
""" op N Binary Indexed Tree
op: (2, )
zero: op (x+e=e+x=x e)
neg : op (1, x+neg(x)=neg(x)+x=e neg(x))
"""
self.op=op
self.zero=zero
self.neg=neg
self.N=N=len(L)
self.log=N.bit_length()-1
X=[zero]*(N+1)
for i in range(N):
p=i+1
X[p]=op(X[p],L[i])
q=p+(p&(-p))
if q<=N:
X[q]=op(X[q], X[p])
self.data=X
def get(self, k):
""" k .
k :
index:
"""
return self.sum(k, k)
def add(self, k, x):
""" k x , .
k :
x :
"""
data=self.data; op=self.op
p=k+1
while p<=self.N:
data[p]=op(self.data[p], x)
p+=p&(-p)
def update(self, k, x):
""" k x , .
k:
x:
"""
a=self.get(k)
y=self.op(self.neg(a), x)
self.add(k,y)
def sum(self, l, r):
""" l r .
※ l != 0 使, .
l:
r:
"""
l=l+1 if 0<=l else 1
r=r+1 if r<self.N else self.N
if l>r:
return self.zero
elif l==1:
return self.__section(r)
else:
return self.op(self.neg(self.__section(l-1)), self.__section(r))
def __section(self, x):
""" B[0]+...+B[x] . """
data=self.data; op=self.op
S=self.zero
while x>0:
S=op(data[x], S)
x-=x&(-x)
return S
def all_sum(self):
return self.sum(0, self.N-1)
def binary_search(self, cond):
""" cond(B[0]+...+B[k]) True k .
cond: 調
※ cond(zero)=True -1 .
※ cond(B[0]+...+B[k]) k (0<=k<N ) N .
"""
if cond(self.zero):
return -1
j=0
r=self.N
t=1<<self.log
data=self.data; op=self.op
alpha=self.zero
while t>0:
if j+t<=self.N:
beta=op(alpha, data[j+t])
if not cond(beta):
alpha=beta
j+=t
t>>=1
return j
def __getitem__(self, index):
if isinstance(index, int):
return self.get(index)
else:
return [self.get(t) for t in index]
def __setitem__(self, index, val):
self.update(index, val)
def __iter__(self):
for k in range(self.N):
yield self.sum(k, k)
#==================================================
from operator import add, neg
from bisect import bisect_left as bis
class Absolute_sum:
def __init__(self, A):
self.A = sorted(A)
self.N = len(A)
self.BIT0 = Binary_Indexed_Tree([0]*self.N, add, 0, neg)
self.BIT1 = Binary_Indexed_Tree([0]*self.N, add, 0, neg)
def get_value(self, i):
p = self.BIT0.binary_search(lambda x:x>=i+1)
return self.BIT1[p]
def insert(self, i):
self.BIT0.add(i, 1)
self.BIT1.add(i, self.A[i])
def remove(self, i):
self.BIT0.add(i, -1)
self.BIT1.add(i, -self.A[i])
def calc(self, black_index, black_value, white_value):
b = self.get_value(black_index)
q = bis(self.A, white_value)
s = self.BIT0.all_sum()
t = self.BIT0.sum(0, q-1)
beta = self.BIT1.sum(q, self.N-1) - self.BIT1.sum(0, q-1) + (2*t-s) *white_value
return beta - abs(b-white_value) + abs(b-black_value)
#==================================================
def solve():
N = int(input())
A = list(map(int, input().split()))
Adj = [[] for _ in range(N+1)]
for j in range(N):
u,v = map(int, input().split())
Adj[u].append(v)
Adj[v].append(u)
root = 0
T = Making_Tree_from_Adjacent_List(N+1, Adj, root, 0)
A_inv = [0] * (N+1)
for i,j in enumerate(sorted(list(range(N+1)), key=lambda i:A[i])):
A_inv[j] = i
Calc = Absolute_sum(sorted(A))
inf = 10 ** 18
Ans = [inf]*(N+1)
for v,c in T.dfs_yielder():
if c == -1:
Calc.remove(A_inv[v])
continue
Calc.insert(A_inv[v])
dep = T.vertex_depth(v)
if dep < 2:
Ans[v] = -1
continue
if dep % 2 == 0:
p = dep // 2 - 1
else:
p = (dep - 1) //2
value_0 = Calc.get_value(0)
value_p = Calc.get_value(p)
value_p1 = Calc.get_value(p+1)
value_last = Calc.get_value(dep)
if value_0 != value_p1:
Ans[v] = min(Ans[v], Calc.calc(0, value_0, value_p1))
if value_last != value_p:
Ans[v] = min(Ans[v], Calc.calc(dep, value_last, value_p))
if value_p == value_p1:
Ans[v] = min(
Ans[v],
Calc.calc(p, value_p+1, value_p),
Calc.calc(p, value_p-1, value_p),
Calc.calc(p, value_p, value_p+1),
Calc.calc(p, value_p, value_p-1))
return Ans[1:]
#==================================================
import sys
input=sys.stdin.readline
write=sys.stdout.write
write("\n".join(map(str, solve())))
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0