#897052 (Python3) No.399 動的な領主

提出ソース

結果

問題	No.399 動的な領主
ユーザー	vwxyz
提出日時	2023-08-01 06:40:59
言語	Python3 (3.12.2 + numpy 1.26.4 + scipy 1.12.0)
結果	AC
実行時間	1,845 ms / 2,000 ms
コード長	11,377 bytes
コンパイル時間	219 ms
コンパイル使用メモリ	14,080 KB
実行使用メモリ	61,388 KB
最終ジャッジ日時	2024-04-19 06:14:52
合計ジャッジ時間	20,841 ms
ジャッジサーバーID （参考情報）	judge5 / judge1

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	44 ms 12,800 KB
testcase_01	AC	44 ms 12,800 KB
testcase_02	AC	45 ms 12,928 KB
testcase_03	AC	45 ms 12,928 KB
testcase_04	AC	54 ms 13,056 KB
testcase_05	AC	169 ms 16,640 KB
testcase_06	AC	1,826 ms 57,092 KB
testcase_07	AC	1,845 ms 56,748 KB
testcase_08	AC	1,815 ms 56,996 KB
testcase_09	AC	1,821 ms 56,132 KB
testcase_10	AC	56 ms 13,056 KB
testcase_11	AC	158 ms 16,768 KB
testcase_12	AC	1,590 ms 55,348 KB
testcase_13	AC	1,620 ms 55,352 KB
testcase_14	AC	1,224 ms 61,388 KB
testcase_15	AC	1,231 ms 61,384 KB
testcase_16	AC	1,413 ms 59,704 KB
testcase_17	AC	1,769 ms 55,196 KB
testcase_18	AC	1,774 ms 56,044 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

import bisect
import copy
import decimal
import fractions
import heapq
import itertools
import math
import random
import sys
import time
from collections import Counter,deque,defaultdict
from functools import lru_cache,reduce
from heapq import heappush,heappop,heapify,heappushpop,_heappop_max,_heapify_max
def _heappush_max(heap,item):
    heap.append(item)
    heapq._siftdown_max(heap, 0, len(heap)-1)
def _heappushpop_max(heap, item):
    if heap and item < heap[0]:
        item, heap[0] = heap[0], item
        heapq._siftup_max(heap, 0)
    return item
from math import gcd as GCD
read=sys.stdin.read
readline=sys.stdin.readline
readlines=sys.stdin.readlines
write=sys.stdout.write

class Graph:
    def __init__(self,V,edges=None,graph=None,directed=False,weighted=False,inf=float("inf")):
        self.V=V
        self.directed=directed
        self.weighted=weighted
        self.inf=inf
        if graph!=None:
            self.graph=graph
            self.edges=[]
            for i in range(self.V):
                if self.weighted:
                    for j,d in self.graph[i]:
                        if self.directed or not self.directed and i<=j:
                            self.edges.append((i,j,d))
                else:
                    for j in self.graph[i]:
                        if self.directed or not self.directed and i<=j:
                            self.edges.append((i,j))
        else:
            self.edges=edges
            self.graph=[[] for i in range(self.V)]
            if weighted:
                for i,j,d in self.edges:
                    self.graph[i].append((j,d))
                    if not self.directed:
                        self.graph[j].append((i,d))
            else:
                for i,j in self.edges:
                    self.graph[i].append(j)
                    if not self.directed:
                        self.graph[j].append(i)

    def SIV_DFS(self,s,bipartite_graph=False,cycle_detection=False,directed_acyclic=False,euler_tour=False,linked_components=False,lowlink=False,parents=False,postorder=False,preorder=False,subtree_size=False,topological_sort=False,unweighted_dist=False,weighted_dist=False):
        seen=[False]*self.V
        finished=[False]*self.V
        if directed_acyclic or cycle_detection or topological_sort:
            dag=True
        if euler_tour:
            et=[]
        if linked_components:
            lc=[]
        if lowlink:
            order=[None]*self.V
            ll=[None]*self.V
            idx=0
        if parents or cycle_detection or lowlink or subtree_size:
            ps=[None]*self.V
        if postorder or topological_sort:
            post=[]
        if preorder:
            pre=[]
        if subtree_size:
            ss=[1]*self.V
        if unweighted_dist or bipartite_graph:
            uwd=[self.inf]*self.V
            uwd[s]=0
        if weighted_dist:
            wd=[self.inf]*self.V
            wd[s]=0
        stack=[(s,0)] if self.weighted else [s]
        while stack:
            if self.weighted:
                x,d=stack.pop()
            else:
                x=stack.pop()
            if not seen[x]:
                seen[x]=True
                stack.append((x,d) if self.weighted else x)
                if euler_tour:
                    et.append(x)
                if linked_components:
                    lc.append(x)
                if lowlink:
                    order[x]=idx
                    ll[x]=idx
                    idx+=1
                if preorder:
                    pre.append(x)
                for y in self.graph[x]:
                    if self.weighted:
                        y,d=y
                    if not seen[y]:
                        stack.append((y,d) if self.weighted else y)
                        if parents or cycle_detection or lowlink or subtree_size:
                            ps[y]=x
                        if unweighted_dist or bipartite_graph:
                            uwd[y]=uwd[x]+1
                        if weighted_dist:
                            wd[y]=wd[x]+d
                    elif not finished[y]:
                        if (directed_acyclic or cycle_detection or topological_sort) and dag:
                            dag=False
                            if cycle_detection:
                                cd=(y,x)
            elif not finished[x]:
                finished[x]=True
                if euler_tour:
                    et.append(~x)
                if lowlink:
                    bl=True
                    for y in self.graph[x]:
                        if self.weighted:
                            y,d=y
                        if ps[x]==y and bl:
                            bl=False
                            continue
                        ll[x]=min(ll[x],order[y])
                    if x!=s:
                        ll[ps[x]]=min(ll[ps[x]],ll[x])
                if postorder or topological_sort:
                    post.append(x)
                if subtree_size:
                    for y in self.graph[x]:
                        if self.weighted:
                            y,d=y
                        if y==ps[x]:
                            continue
                        ss[x]+=ss[y]
        if bipartite_graph:
            bg=[[],[]]
            for tpl in self.edges:
                x,y=tpl[:2] if self.weighted else tpl
                if uwd[x]==self.inf or uwd[y]==self.inf:
                    continue
                if not uwd[x]%2^uwd[y]%2:
                    bg=False
                    break
            else:
                for x in range(self.V):
                    if uwd[x]==self.inf:
                        continue
                    bg[uwd[x]%2].append(x)
        retu=()
        if bipartite_graph:
            retu+=(bg,)
        if cycle_detection:
            if dag:
                cd=[]
            else:
                y,x=cd
                cd=self.Route_Restoration(y,x,ps)
            retu+=(cd,)
        if directed_acyclic:
            retu+=(dag,)
        if euler_tour:
            retu+=(et,)
        if linked_components:
            retu+=(lc,)
        if lowlink:
            retu=(ll,)
        if parents:
            retu+=(ps,)
        if postorder:
            retu+=(post,)
        if preorder:
            retu+=(pre,)
        if subtree_size:
            retu+=(ss,)
        if topological_sort:
            if dag:
                tp_sort=post[::-1]
            else:
                tp_sort=[]
            retu+=(tp_sort,)
        if unweighted_dist:
            retu+=(uwd,)
        if weighted_dist:
            retu+=(wd,)
        if len(retu)==1:
            retu=retu[0]
        return retu

    def Build_LCA(self,s,segment_tree=False):
        self.lca_segment_tree=segment_tree
        if self.lca_segment_tree:
            self.lca_euler_tour,self.lca_parents,depth=self.SIV_DFS(s,euler_tour=True,parents=True,unweighted_dist=True)
            self.lca_dfs_in_index=[None]*self.V
            self.lca_dfs_out_index=[None]*self.V
            for i,x in enumerate(self.lca_euler_tour):
                if x>=0:
                    self.lca_dfs_in_index[x]=i
                else:
                    self.lca_dfs_out_index[~x]=i
            self.ST=Segment_Tree(2*self.V,lambda x,y:min(x,y),self.V)
            lst=[None]*(2*self.V)
            for i in range(2*self.V-1):
                if self.lca_euler_tour[i]>=0:
                    lst[i]=depth[self.lca_euler_tour[i]]
                else:
                    lst[i]=depth[self.lca_parents[~self.lca_euler_tour[i]]]
            lst[2*self.V-1]=-1
            self.ST.Build(lst)
        else:
            self.lca_parents,self.lca_depth=self.SIV_DFS(s,parents=True,unweighted_dist=True)
            self.lca_PD=Path_Doubling(self.V,self.lca_parents)
            self.lca_PD.Build_Next(self.V)

    def LCA(self,a,b):
        if self.lca_segment_tree:
            m=min(self.lca_dfs_in_index[a],self.lca_dfs_in_index[b])
            M=max(self.lca_dfs_in_index[a],self.lca_dfs_in_index[b])
            x=self.lca_euler_tour[self.ST.Fold_Index(m,M+1)]
            if x>=0:
                lca=x
            else:
                lca=self.lca_parents[~x]
        else:
            if self.lca_depth[a]>self.lca_depth[b]:
                a,b=b,a
            b=self.lca_PD.Permutation_Doubling(b,self.lca_depth[b]-self.lca_depth[a])
            if a!=b:
                for k in range(self.lca_PD.k-1,-1,-1):
                    if self.lca_PD.permutation_doubling[k][a]!=self.lca_PD.permutation_doubling[k][b]:
                        a,b=self.lca_PD.permutation_doubling[k][a],self.lca_PD.permutation_doubling[k][b]
                a,b=self.lca_PD.permutation_doubling[0][a],self.lca_PD.permutation_doubling[0][b]
            lca=a
        return lca

class Path_Doubling:
    def __init__(self,N,permutation,lst=None,f=None,e=None):
        self.N=N
        self.permutation=permutation
        self.lst=lst
        self.f=f
        self.e=e

    def Build_Next(self,K=None):
        if K==None:
            K=self.N
        self.k=K.bit_length()
        self.permutation_doubling=[[None]*self.N for k in range(self.k)]
        for n in range(self.N):
            self.permutation_doubling[0][n]=self.permutation[n]
        if self.lst!=None:
            self.doubling=[[None]*self.N for k in range(self.k)]
            for n in range(self.N):
                self.doubling[0][n]=self.lst[n]
        for k in range(1,self.k):
            for n in range(self.N):
                if self.permutation_doubling[k-1][n]!=None:
                    self.permutation_doubling[k][n]=self.permutation_doubling[k-1][self.permutation_doubling[k-1][n]]
                    if self.f!=None:
                        self.doubling[k][n]=self.f(self.doubling[n][k-1],self.doubling[k-1][self.permutation_doubling[k-1][n]])

    def Permutation_Doubling(self,N,K):
        if K<0 or 1<<self.k<=K:
            return None
        for k in range(self.k):
            if K>>k&1 and N!=None:
                N=self.permutation_doubling[k][N]
        return N

    def Doubling(self,N,K):
        if K<0:
            return self.e
        retu=self.e
        for k in range(self.k):
            if K>>k&1:
                if self.permutation_doubling[k][N]==None:
                    return None
                retu=self.f(retu,self.doubling[k][N])
                N=self.permutation_doubling[k][N]
        return N,retu

    def Bisect(self,x,is_ok):
        if not is_ok(x):
            return -1,None
        K=0
        for k in range(self.k-1,-1,-1):
            if K|1<<k<=self.K and is_ok(self.permutation_doubling[k][x]):
                K|=1<<k
                x=self.permutation_doubling[k][x]
        return K,x

N=int(readline())
edges=[]
for n in range(N-1):
    u,v=map(int,readline().split())
    u-=1;v-=1
    edges.append((u,v))
G=Graph(N,edges=edges)
G.Build_LCA(0)
parents,tour=G.SIV_DFS(0,parents=True,postorder=True)
imos=[0]*N
Q=int(readline())
for q in range(Q):
    A,B=map(int,readline().split())
    A-=1;B-=1
    imos[A]+=1
    imos[B]+=1
    lca=G.LCA(A,B)
    imos[lca]-=1
    if parents[lca]!=None:
        imos[parents[lca]]-=1
for x in tour:
    for y in G.graph[x]:
        if y==parents[x]:
            continue
        imos[x]+=imos[y]
ans=sum(c*(c+1)//2 for c in imos)
print(ans)

yukicoder

結果

テストケース

ソースコード