#709605 (PyPy3) No.1718 Random Squirrel

提出ソース
結果

問題	No.1718 Random Squirrel
ユーザー	chineristAC
提出日時	2021-10-22 22:43:53
言語	PyPy3 (7.3.15)
結果	AC
実行時間	449 ms / 2,000 ms
コード長	17,582 bytes
コンパイル時間	228 ms
コンパイル使用メモリ	81,876 KB
実行使用メモリ	119,488 KB
最終ジャッジ日時	2023-10-24 14:06:05
合計ジャッジ時間	10,081 ms
ジャッジサーバーID （参考情報）	judge15 / judge11
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	51 ms 60,404 KB
testcase_01	AC	46 ms 60,404 KB
testcase_02	AC	46 ms 60,404 KB
testcase_03	AC	47 ms 60,404 KB
testcase_04	AC	48 ms 60,404 KB
testcase_05	AC	47 ms 60,404 KB
testcase_06	AC	47 ms 60,404 KB
testcase_07	AC	48 ms 60,404 KB
testcase_08	AC	47 ms 60,404 KB
testcase_09	AC	47 ms 60,404 KB
testcase_10	AC	46 ms 60,404 KB
testcase_11	AC	312 ms 98,212 KB
testcase_12	AC	226 ms 83,656 KB
testcase_13	AC	334 ms 94,376 KB
testcase_14	AC	333 ms 99,384 KB
testcase_15	AC	330 ms 103,280 KB
testcase_16	AC	258 ms 88,464 KB
testcase_17	AC	312 ms 96,820 KB
testcase_18	AC	408 ms 118,760 KB
testcase_19	AC	346 ms 102,468 KB
testcase_20	AC	218 ms 85,268 KB
testcase_21	AC	394 ms 106,004 KB
testcase_22	AC	443 ms 119,488 KB
testcase_23	AC	447 ms 106,192 KB
testcase_24	AC	404 ms 106,776 KB
testcase_25	AC	449 ms 119,424 KB
testcase_26	AC	220 ms 105,424 KB
testcase_27	AC	235 ms 105,196 KB
testcase_28	AC	209 ms 105,152 KB
testcase_29	AC	250 ms 108,276 KB
testcase_30	AC	257 ms 115,368 KB
testcase_31	AC	224 ms 107,604 KB
testcase_32	AC	249 ms 108,268 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
def divisors(M):
    d=[]
    i=1
    while M>=i**2:
        if M%i==0:
            d.append(i)
            if i**2!=M:
                d.append(M//i)
        i=i+1
    return d

def popcount(x):
    x = x - ((x >> 1) & 0x55555555)
    x = (x & 0x33333333) + ((x >> 2) & 0x33333333)
    x = (x + (x >> 4)) & 0x0f0f0f0f
    x = x + (x >> 8)
    x = x + (x >> 16)
    return x & 0x0000007f

def eratosthenes(n):
    res=[0 for i in range(n+1)]
    prime=set([])
    for i in range(2,n+1):
        if not res[i]:
            prime.add(i)
            for j in range(1,n//i+1):
                res[i*j]=1
    return prime

def factorization(n):
    res=[]
    for p in prime:
        if n%p==0:
            while n%p==0:
                n//=p
            res.append(p)
    if n!=1:
        res.append(n)
    return res

def euler_phi(n):
    res = n
    for x in range(2,n+1):
        if x ** 2 > n:
            break
        if n%x==0:
            res = res//x * (x-1)
            while n%x==0:
                n //= x
    if n!=1:
        res = res//n * (n-1)
    return res

def ind(b,n):
    res=0
    while n%b==0:
        res+=1
        n//=b
    return res

def isPrimeMR(n):
    d = n - 1
    d = d // (d & -d)
    L = [2, 3, 5, 7, 11, 13, 17]
    for a in L:
        t = d
        y = pow(a, t, n)
        if y == 1: continue
        while y != n - 1:
            y = (y * y) % n
            if y == 1 or t == n - 1: return 0
            t <<= 1
    return 1
def findFactorRho(n):
    from math import gcd
    m = 1 << n.bit_length() // 8
    for c in range(1, 99):
        f = lambda x: (x * x + c) % n
        y, r, q, g = 2, 1, 1, 1
        while g == 1:
            x = y
            for i in range(r):
                y = f(y)
            k = 0
            while k < r and g == 1:
                ys = y
                for i in range(min(m, r - k)):
                    y = f(y)
                    q = q * abs(x - y) % n
                g = gcd(q, n)
                k += m
            r <<= 1
        if g == n:
            g = 1
            while g == 1:
                ys = f(ys)
                g = gcd(abs(x - ys), n)
        if g < n:
            if isPrimeMR(g): return g
            elif isPrimeMR(n // g): return n // g
            return findFactorRho(g)
def primeFactor(n):
    i = 2
    ret = {}
    rhoFlg = 0
    while i*i <= n:
        k = 0
        while n % i == 0:
            n //= i
            k += 1
        if k: ret[i] = k
        i += 1 + i % 2
        if i == 101 and n >= 2 ** 20:
            while n > 1:
                if isPrimeMR(n):
                    ret[n], n = 1, 1
                else:
                    rhoFlg = 1
                    j = findFactorRho(n)
                    k = 0
                    while n % j == 0:
                        n //= j
                        k += 1
                    ret[j] = k

    if n > 1: ret[n] = 1
    if rhoFlg: ret = {x: ret[x] for x in sorted(ret)}
    return ret

def divisors(n):
    res = [1]
    prime = primeFactor(n)
    for p in prime:
        newres = []
        for d in res:
            for j in range(prime[p]+1):
                newres.append(d*p**j)
        res = newres
    res.sort()
    return res

def xorconv(n,X,Y):
    if n==0:
        res=[(X[0]*Y[0])%mod]
        return res
    x=[X[i]+X[i+2**(n-1)] for i in range(2**(n-1))]
    y=[Y[i]+Y[i+2**(n-1)] for i in range(2**(n-1))]
    z=[X[i]-X[i+2**(n-1)] for i in range(2**(n-1))]
    w=[Y[i]-Y[i+2**(n-1)] for i in range(2**(n-1))]
    res1=xorconv(n-1,x,y)
    res2=xorconv(n-1,z,w)
    former=[(res1[i]+res2[i])*inv for i in range(2**(n-1))]
    latter=[(res1[i]-res2[i])*inv for i in range(2**(n-1))]
    former=list(map(lambda x:x%mod,former))
    latter=list(map(lambda x:x%mod,latter))
    return former+latter

def merge_sort(A,B):
    pos_A,pos_B = 0,0
    n,m = len(A),len(B)
    res = []
    while pos_A < n and pos_B < m:
        a,b = A[pos_A],B[pos_B]
        if a < b:
            res.append(a)
            pos_A += 1
        else:
            res.append(b)
            pos_B += 1
    res += A[pos_A:]
    res += B[pos_B:]
    return res

class UnionFindVerSize():
    def __init__(self, N):
        self._parent = [n for n in range(0, N)]
        self._size = [1] * N
        self.group = N

    def find_root(self, x):
        if self._parent[x] == x: return x
        self._parent[x] = self.find_root(self._parent[x])
        stack = [x]
        while self._parent[stack[-1]]!=stack[-1]:
            stack.append(self._parent[stack[-1]])
        for v in stack:
            self._parent[v] = stack[-1]
        return self._parent[x]

    def unite(self, x, y):
        gx = self.find_root(x)
        gy = self.find_root(y)
        if gx == gy: return

        self.group -= 1

        if self._size[gx] < self._size[gy]:
            self._parent[gx] = gy
            self._size[gy] += self._size[gx]
        else:
            self._parent[gy] = gx
            self._size[gx] += self._size[gy]

    def get_size(self, x):
        return self._size[self.find_root(x)]

    def is_same_group(self, x, y):
        return self.find_root(x) == self.find_root(y)

class WeightedUnionFind():
    def __init__(self,N):
        self.parent = [i for i in range(N)]
        self.size = [1 for i in range(N)]
        self.val = [0 for i in range(N)]
        self.flag = True
        self.edge = [[] for i in range(N)]

    def dfs(self,v,pv):
        stack = [(v,pv)]
        new_parent = self.parent[pv]
        while stack:
            v,pv = stack.pop()
            self.parent[v] = new_parent
            for nv,w in self.edge[v]:
                if nv!=pv:
                    self.val[nv] = self.val[v] + w
                    stack.append((nv,v))

    def unite(self,x,y,w):
        if not self.flag:
            return
        if self.parent[x]==self.parent[y]:
            self.flag = (self.val[x] - self.val[y] == w)
            return

        if self.size[self.parent[x]]>self.size[self.parent[y]]:
            self.edge[x].append((y,-w))
            self.edge[y].append((x,w))
            self.size[x] += self.size[y]
            self.val[y] = self.val[x] - w
            self.dfs(y,x)
        else:
            self.edge[x].append((y,-w))
            self.edge[y].append((x,w))
            self.size[y] += self.size[x]
            self.val[x] = self.val[y] + w
            self.dfs(x,y)

class Dijkstra():
    class Edge():
        def __init__(self, _to, _cost):
            self.to = _to
            self.cost = _cost

    def __init__(self, V):
        self.G = [[] for i in range(V)]
        self._E = 0
        self._V = V

    @property
    def E(self):
        return self._E

    @property
    def V(self):
        return self._V

    def add_edge(self, _from, _to, _cost):
        self.G[_from].append(self.Edge(_to, _cost))
        self._E += 1

    def shortest_path(self, s):
        import heapq
        que = []
        d = [10**15] * self.V
        d[s] = 0
        heapq.heappush(que, (0, s))

        while len(que) != 0:
            cost, v = heapq.heappop(que)
            if d[v] < cost: continue

            for i in range(len(self.G[v])):
                e = self.G[v][i]
                if d[e.to] > d[v] + e.cost:
                    d[e.to] = d[v] + e.cost
                    heapq.heappush(que, (d[e.to], e.to))
        return d

#Z[i]:length of the longest list starting from S[i] which is also a prefix of S
#O(|S|)
def Z_algorithm(s):
    N = len(s)
    Z_alg = [0]*N

    Z_alg[0] = N
    i = 1
    j = 0
    while i < N:
        while i+j < N and s[j] == s[i+j]:
            j += 1
        Z_alg[i] = j
        if j == 0:
            i += 1
            continue
        k = 1
        while i+k < N and k + Z_alg[k]<j:
            Z_alg[i+k] = Z_alg[k]
            k += 1
        i += k
        j -= k
    return Z_alg

class BIT():
    def __init__(self,n,mod=0):
        self.BIT = [0]*(n+1)
        self.num = n
        self.mod = mod

    def query(self,idx):
        res_sum = 0
        mod = self.mod
        while idx > 0:
            res_sum += self.BIT[idx]
            if mod:
                res_sum %= mod
            idx -= idx&(-idx)
        return res_sum

    #Ai += x O(logN)
    def update(self,idx,x):
        mod = self.mod
        while idx <= self.num:
            self.BIT[idx] += x
            if mod:
                self.BIT[idx] %= mod
            idx += idx&(-idx)
        return

class dancinglink():
    def __init__(self,n,debug=False):
        self.n = n
        self.debug = debug
        self._left = [i-1 for i in range(n)]
        self._right = [i+1 for i in range(n)]
        self.exist = [True for i in range(n)]

    def pop(self,k):
        if self.debug:
            assert self.exist[k]
        L = self._left[k]
        R = self._right[k]
        if L!=-1:
            if R!=self.n:
                self._right[L],self._left[R] = R,L
            else:
                self._right[L] = self.n
        elif R!=self.n:
            self._left[R] = -1
        self.exist[k] = False

    def left(self,idx,k=1):
        if self.debug:
            assert self.exist[idx]
        res = idx
        while k:
            res = self._left[res]
            if res==-1:
                break
            k -= 1
        return res

    def right(self,idx,k=1):
        if self.debug:
            assert self.exist[idx]
        res = idx
        while k:
            res = self._right[res]
            if res==self.n:
                break
            k -= 1
        return res

class SparseTable():
    def __init__(self,A,merge_func,ide_ele):
        N = len(A)
 
        self.merge_func = merge_func
 
        self.lg = [0]*(N + 1)
        for i in range(2, N+1):
            self.lg[i] = self.lg[i >> 1] + 1
        self.pow_2 = [pow(2,i) for i in range(20)]
 
        self.table = [None]*(self.lg[N] + 1)
        st0 = self.table[0] = [a for a in A]
        b = 1
        for i in range(self.lg[N]):
            st0 = self.table[i+1] = [self.merge_func(u,v) for u, v in zip(st0, st0[b:])]
            b <<= 1
 
    def query(self,s,t):
        b = t-s+1
        m = self.lg[b]
        return self.merge_func(self.table[m][s],self.table[m][t-self.pow_2[m]+1])
 

class BinaryTrie:
    class node:
        def __init__(self,val):
            self.left = None
            self.right = None
            self.max = val

    def __init__(self):
        self.root = self.node(-10**15)

    def append(self,key,val):
        pos = self.root
        for i in range(29,-1,-1):
            pos.max = max(pos.max,val)
            if key>>i & 1:
                if pos.right is None:
                    pos.right = self.node(val)
                    pos = pos.right
                else:
                    pos = pos.right
            else:
                if pos.left is None:
                    pos.left = self.node(val)
                    pos = pos.left
                else:
                    pos = pos.left
        pos.max = max(pos.max,val)

    def search(self,M,xor):
        res = -10**15
        pos = self.root
        for i in range(29,-1,-1):
            if pos is None:
                break

            if M>>i & 1:
                if xor>>i & 1:
                    if pos.right:
                        res = max(res,pos.right.max)
                    pos = pos.left
                else:
                    if pos.left:
                        res = max(res,pos.left.max)
                    pos = pos.right
            else:
                if xor>>i & 1:
                    pos = pos.right
                else:
                    pos = pos.left

        if pos:
            res = max(res,pos.max)
        return res

def solveequation(edge,ans,n,m):
    #edge=[[to,dire,id]...]
    def dfs(v):
        used[v]=True
        r=ans[v]
        for to,dire,id in edge[v]:
            if used[to]:
                continue
            y=dfs(to)
            if dire==-1:
                x[id]=y
            else:
                x[id]=-y
            r+=y
        return r
        
    x=[0]*m
    used=[False]*n
    for v in range(n):
        if used[v]:
            continue
        y = dfs(v)
        if y!=0:
            return False
    return x

    

class slope_trick():
    def __init__(self):
        self.L = [10**17]
        self.R = [10**17]
        self.min_f = 0

        self.x_left = 0
        self.x_right = 0

    def add_right(self,a):
        a -= self.x_left
        l0 = -self.L[0]
        self.min_f  = self.min_f + max(0,l0-a)
        if l0 <= a:
            a += self.x_left
            a -= self.x_right
            heappush(self.R,a)
        else:
            heappush(self.L,-a)
            a = -heappop(self.L)
            a += self.x_left
            a -= self.x_right
            heappush(self.R,a)

        #self.min_f  = self.min_f + max(0,l0-a)

    def add_left(self,a):
        a -= self.x_right
        r0 = self.R[0]
        self.min_f = self.min_f + max(0,a-r0)

        if a <= r0:
            a += self.x_right
            a -= self.x_left
            heappush(self.L,-a)
        else:
            heappush(self.R,a)
            a = heappop(self.R)
            a += self.x_right
            a -= self.x_left
            heappush(self.L,-a)

        #self.min_f = self.min_f + max(0,a-r0)

    def add_abs(self,a):
        self.add_left(a)
        self.add_right(a)

    def change_min_slide(self,a,b):
        self.x_left += a
        self.x_right += b
    
    def get_val(self,x):
        L = [-l+self.x_left for l in self.L]
        L.sort()
        R = [r+self.x_right for r in self.R]
        R.sort()

        res = self.min_f

        if 0 < L[-1]:
            L = L[::-1]
            n = len(L)
            for i in range(n):
                c0 = L[i]
                c1 = L[i+1]
        
                if c1 <= x <= c0:
                    res += (i+1) * (c0-x)
                    break
                else:
                    res += (i+1) * (c0-c1)
            return res
        elif L[-1] <= x <= R[0]:
            return res
        else:
            n = len(R)
            for i in range(n):
                c0 = R[i]
                c1 = R[i+1]
                if c0 <= x <= c1:
                    res += (i+1) * (x-c0)
                    break
                else:
                    res += (i+1) * (c1-c0)
            return res

def fwt(n,A):
    assert len(A) == 2**n
    for i in range(n):
        t = 2**i
        for j in range(2**n):
            if j&t==0:
                A[j] += A[j|t]
    return A

def ifwt(n,A):
    assert len(A) == 2**n
    for i in range(n):
        t = 2**i
        for j in range(2**n):
            if j&t==0:
                A[j] -= A[j|t]
    return A

import sys,random
from collections import deque

input = lambda :sys.stdin.readline().rstrip()
mi = lambda :map(int,input().split())
li = lambda :list(mi())

N,K = mi()
edge = [[] for i in range(N)]
for _ in range(N-1):
    u,v = mi()
    edge[u-1].append(v-1)
    edge[v-1].append(u-1)

D = li()
D = set([d-1 for d in D])

ans = [0] * N

parent = [-1] * N
deq = deque([0])
topo = []
while deq:
    v = deq.popleft()
    topo.append(v)
    for nv in edge[v]:
        if nv!=parent[v]:
            parent[nv] = v
            deq.append(nv)

for v in range(N):
    edge[v] = [nv for nv in edge[v] if nv!=parent[v]]

bottom_size = [0]*N
bottom_M = [-1]*N
bottom_ds = [0]*N
for v in topo[::-1]:
    for nv in edge[v]:
        if nv!=parent[v]:
            bottom_ds[v] += bottom_size[nv] + bottom_ds[nv]
            bottom_size[v] += bottom_size[nv]
            if bottom_size[nv]:
                bottom_M[v] = max(bottom_M[v],bottom_M[nv]+1)
    if v in D:
        bottom_size[v] += 1
        bottom_M[v] = max(bottom_M[v],0)


top_size = [0]*N
top_ds = [0]*N
top_M = [-1]*N
for v in topo:
    for nv in edge[v]:
        top_size[nv] = bottom_size[v] - bottom_size[nv] + top_size[v]
        top_ds[nv] = bottom_ds[v] - (bottom_ds[nv]+bottom_size[nv]) +  top_size[nv] + top_ds[v]
        if top_M[v]!=-1:
            top_M[nv] = max(top_M[nv],top_M[v]+1)
        if v in D:
            top_M[nv] = max(top_M[nv],1)
    
    left = [bottom_M[nv] for nv in edge[v]]
    right = [bottom_M[nv] for nv in edge[v]]
    n = len(edge[v])
    for i in range(1,n):
        left[i] = max(left[i],left[i-1])
    for i in range(n-2,-1,-1):
        right[i] = max(right[i],right[i+1])
    
    for i in range(n):
        tmp = -1
        if i!=0 and left[i-1]!=-1:
            tmp = max(tmp,left[i-1]+2)
        if i!=n-1 and right[i+1]!=-1:
            tmp = max(tmp,right[i+1]+2)

        nv = edge[v][i]
        top_M[nv] = max(top_M[nv],tmp)
    
is_tree = [False for v in range(N)]
for v in range(N):
    if v in D:
        is_tree[v] = True
    else:
        if top_size[v]:
            if bottom_size[v]:
                is_tree[v] = True
        else:
            cnt = 0
            for nv in edge[v]:
                if bottom_size[nv]:
                    cnt += 1
            if cnt >= 2:
                is_tree[v] = True

S = sum(is_tree)
dist = [N+1 for v in range(N)]
for v in range(N):
    if is_tree[v]:
        dist[v] = 0
deq = deque([v for v in range(N) if is_tree[v]])
while deq:
    v = deq.popleft()
    for nv in edge[v]:
        if dist[nv] > dist[v] + 1:
            dist[nv] = dist[v] + 1
            deq.append(nv)
    nv = parent[v]
    if nv!=-1 and dist[nv] > dist[v] + 1:
        dist[nv] = dist[v] + 1
        deq.append(nv)

for v in range(N):
    ans[v] = 2 * (S-1) + 2 * dist[v] - max(bottom_M[v],top_M[v])
    print(ans[v])
yukicoder

結果

テストケース

ソースコード
