結果
問題 | No.2654 [Cherry 6th Tune] Re: start! (Black Sheep) |
ユーザー |
👑 ![]() |
提出日時 | 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 |
ソースコード
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=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):""" 頂点 root を根に設定する."""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 を設定する (頂点 x の方が親)."""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):""" 木の深さを求める.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.parentfor x in range(self.index,self.index+self.N):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 range(self.index,self.index+self.N):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_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,dxx,y=y,xpa=self.parentwhile dx>dy:x=pa[x]dx-=1while x!=y:x=pa[x]y=pa[y]return xdef __lca_prepare(self):assert self.__after_seal_check()N=self.Nbit=max(1, ((2*N-1)-1).bit_length())D=[[0]*(2*N-1) for _ in range(bit)]self.euler_tour_vertex()tour=self.euler_vertexD[0]=tour.copy()dep=self.depth_search(True)for i in range(1, bit):shift=1<<itab=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:breaktab[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=Dreturndef 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,xa,b=b,aif a==b:return self.lca_dst[0][a]p=(a^b).bit_length()-1tab=self.lca_dst[p]u=tab[a]; v=tab[b]return u if self.vertex_depth(u)<self.vertex_depth(v) else vdef 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])+1def 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, 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.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_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 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 direction(self, u, v):""" 頂点 u から頂点 v (u!=v) へ向かうパスが頂点 u の次に通る頂点"""assert self.__after_seal_check(u,v)assert u!=vif 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: intv: intk: intdefault=-1: int"""assert self.__after_seal_check(u,v)if k==0:return u# lca を求める.x=u; y=vdx=self.vertex_depth(x); dy=self.vertex_depth(y)if dx>dy:x,y=y,xdx,dy=dy,dxy=self.upper(y, dy-dx)if x==self.root or x==y:w=xelse:bit=dx.bit_length()X=self.upper_listfor t in range(bit-1,-1,-1):px=X[t][x]; py=X[t][y]if px!=py:x=px; y=pyw=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 defaultelif 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_depthif 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"):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, 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.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)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: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,..., 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.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]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)※ 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.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_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.rootch=self.childrenif order!=None:for i in range(self.index,self.index+self.N):ch[i].sort(key=order)pa=self.parentR=[-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]=vif R[v]==S[v]:v=pa[v]else: #進めるw=vv=ch[v][S[v]]S[w]+=1self.euler_vertex=Xself.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]=telse:self.out_time[v]=tdef euler_tour_edge(self):""" オイラーツアー (edge) に関する計算を行う.(u, v, k): u から v へ向かう (k=+1 のときは葉へ進む向き, k=-1 のときは根へ進む向き)"""assert self.__after_seal_check()if hasattr(self,"euler_edge"):returnif not hasattr(self, "euler_vertex"):self.euler_tour_vertex()self.euler_edge=[0]*(2*(self.N-1))euler=self.euler_vertexpa=self.parentfor t in range(2*(self.N-1)):u=euler[t]; v=euler[t+1]k=1 if u==pa[v] else -1self.euler_edge[t]=(u,v,k)def 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 Gdef 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=0for i in range(K-1):X+=self.distance(S[i],S[i+1])return (X+self.distance(S[-1],S[0]))//2def Making_Tree_from_Adjacent_List(N, A, root, index=0):""" 隣接リストから木を作る."""from collections import dequeT=Tree(N, index)T.root_set(root)S=[False]*(N+index); S[root]=TrueQ=deque([root])while Q:v=Q.popleft()for w in A[v]:if not S[w]:S[w]=TrueT.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=opself.zero=zeroself.neg=negself.N=N=len(L)self.log=N.bit_length()-1X=[zero]*(N+1)for i in range(N):p=i+1X[p]=op(X[p],L[i])q=p+(p&(-p))if q<=N:X[q]=op(X[q], X[p])self.data=Xdef 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.opp=k+1while 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 1r=r+1 if r<self.N else self.Nif l>r:return self.zeroelif 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.opS=self.zerowhile x>0:S=op(data[x], S)x-=x&(-x)return Sdef 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 -1j=0r=self.Nt=1<<self.logdata=self.data; op=self.opalpha=self.zerowhile t>0:if j+t<=self.N:beta=op(alpha, data[j+t])if not cond(beta):alpha=betaj+=tt>>=1return jdef __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, negfrom bisect import bisect_left as bisclass 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_valuereturn 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 = 0T = 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] = iCalc = Absolute_sum(sorted(A))inf = 10 ** 18Ans = [inf]*(N+1)for v,c in T.dfs_yielder():if c == -1:Calc.remove(A_inv[v])continueCalc.insert(A_inv[v])dep = T.vertex_depth(v)if dep < 2:Ans[v] = -1continueif dep % 2 == 0:p = dep // 2 - 1else:p = (dep - 1) //2value_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 sysinput=sys.stdin.readlinewrite=sys.stdout.writewrite("\n".join(map(str, solve())))