結果

問題 No.2892 Lime and Karin
ユーザー chineristACchineristAC
提出日時 2024-09-13 21:52:30
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 2,566 ms / 8,000 ms
コード長 4,503 bytes
コンパイル時間 267 ms
コンパイル使用メモリ 81,792 KB
実行使用メモリ 171,384 KB
最終ジャッジ日時 2024-09-13 21:54:06
合計ジャッジ時間 65,692 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge5
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 48 ms
56,448 KB
testcase_01 AC 48 ms
56,576 KB
testcase_02 AC 53 ms
56,524 KB
testcase_03 AC 49 ms
56,192 KB
testcase_04 AC 51 ms
56,576 KB
testcase_05 AC 55 ms
56,960 KB
testcase_06 AC 51 ms
56,960 KB
testcase_07 AC 48 ms
56,576 KB
testcase_08 AC 55 ms
63,104 KB
testcase_09 AC 55 ms
63,104 KB
testcase_10 AC 51 ms
56,960 KB
testcase_11 AC 49 ms
56,320 KB
testcase_12 AC 49 ms
57,088 KB
testcase_13 AC 51 ms
57,472 KB
testcase_14 AC 1,054 ms
111,684 KB
testcase_15 AC 991 ms
110,796 KB
testcase_16 AC 1,742 ms
132,168 KB
testcase_17 AC 456 ms
91,172 KB
testcase_18 AC 1,582 ms
133,084 KB
testcase_19 AC 2,308 ms
152,000 KB
testcase_20 AC 270 ms
82,096 KB
testcase_21 AC 1,196 ms
118,396 KB
testcase_22 AC 1,521 ms
130,368 KB
testcase_23 AC 678 ms
100,300 KB
testcase_24 AC 2,432 ms
154,460 KB
testcase_25 AC 2,221 ms
155,148 KB
testcase_26 AC 2,288 ms
153,672 KB
testcase_27 AC 2,252 ms
153,440 KB
testcase_28 AC 2,354 ms
154,532 KB
testcase_29 AC 2,566 ms
153,864 KB
testcase_30 AC 2,242 ms
156,908 KB
testcase_31 AC 2,370 ms
153,084 KB
testcase_32 AC 2,409 ms
153,844 KB
testcase_33 AC 2,348 ms
152,772 KB
testcase_34 AC 455 ms
126,852 KB
testcase_35 AC 475 ms
127,240 KB
testcase_36 AC 457 ms
126,508 KB
testcase_37 AC 457 ms
127,428 KB
testcase_38 AC 488 ms
127,620 KB
testcase_39 AC 2,151 ms
171,060 KB
testcase_40 AC 2,295 ms
170,784 KB
testcase_41 AC 2,318 ms
170,976 KB
testcase_42 AC 2,351 ms
170,712 KB
testcase_43 AC 2,284 ms
171,384 KB
testcase_44 AC 765 ms
105,160 KB
testcase_45 AC 1,943 ms
144,416 KB
testcase_46 AC 936 ms
109,216 KB
testcase_47 AC 1,193 ms
121,524 KB
testcase_48 AC 529 ms
95,452 KB
testcase_49 AC 1,434 ms
127,864 KB
testcase_50 AC 1,236 ms
139,408 KB
testcase_51 AC 1,268 ms
139,732 KB
testcase_52 AC 1,274 ms
138,748 KB
testcase_53 AC 1,309 ms
139,252 KB
testcase_54 AC 1,311 ms
139,168 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
from itertools import permutations
from heapq import heappop,heappush
from collections import deque
import random
import bisect

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

def solve(N,edge,S):

    centroid_done = [False] * N

    tmp_parent = [None] * N
    tmp_sz = [None] * N

    def find_centroid(root):
        deq = deque([root])
        topo = []
        while deq:
            v = deq.popleft()
            topo.append(v)
            for nv in edge[v]:
                if centroid_done[nv]:
                    continue
                if nv == tmp_parent[v]:
                    continue
                tmp_parent[nv] = v
                deq.append(nv)

        tmp_n = len(topo)
        centroid = -1
        for v in topo[::-1]:
            tmp_sz[v] = 1
            centroid_flg = True
            for nv in edge[v]:
                if centroid_done[nv]:
                    continue
                if nv == tmp_parent[v]:
                    continue
                
                if tmp_sz[nv]*2 > tmp_n:
                    centroid_flg = False
                tmp_sz[v] += tmp_sz[nv]
            
            if 2*(tmp_n-tmp_sz[v]) > tmp_n:
                centroid_flg = False
            if centroid_flg:
                centroid = v
        
        for v in topo:
            tmp_sz[v] = None
            tmp_parent[v] = None
        
        return centroid
    
    def f(n,A,x):
        """
        A[i]+A[j]+x > 0 を満たすi,jの数を求める
        -n <= A[i] <= n が保証される
        """

        tmp_cnt = [0] * (2*n+1)
        for a in A:
            tmp_cnt[a+n] += 1
        for i in range(2*n)[::-1]:
            tmp_cnt[i] += tmp_cnt[i+1]
        
        res = 0
        for a in A:
            lower = max((-a-x+1) + n,0)
            if lower <= 2*n:
                res += tmp_cnt[lower]
            if a+a+x > 0:
                res -= 1
        return res//2
    
    tmp_dep = [None] * N
    tmp_centroid_child_parent = [None] * N
    tmp_centroid_child_group = [[] for v in range(N)]
                
    def calc_sub(centroid):
        deq = deque([centroid])
        tmp_dep[centroid] = 0
        topo = []
        while deq:
            v = deq.popleft()
            topo.append(v)
            for nv in edge[v]:
                if centroid_done[nv]:
                    continue
                if nv == tmp_parent[v]:
                    continue
                tmp_parent[nv] = v

                if v == centroid:
                    tmp_centroid_child_parent[nv] = nv
                else:
                    tmp_centroid_child_parent[nv] = tmp_centroid_child_parent[v]
                tmp_centroid_child_group[tmp_centroid_child_parent[nv]].append(nv)

                tmp_dep[nv] = tmp_dep[v] + 2 * S[nv] - 1
                deq.append(nv)
        
        tmp_n = len(topo)
        A = [tmp_dep[v] for v in topo if v!=centroid]
        
        res = f(tmp_n,A,2 * S[centroid] - 1)
        for v in topo:
            if v != centroid and tmp_dep[v] + 2 * S[centroid] - 1 > 0:
                res += 1
        

        for centroid_child in edge[centroid]:
            if centroid_done[centroid_child]:
                continue
            
            tmp_n = len(tmp_centroid_child_group[centroid_child])
            A = [tmp_dep[v] for v in tmp_centroid_child_group[centroid_child]]
            res -= f(tmp_n,A,2 * S[centroid] - 1)
        
        for v in topo:
            tmp_parent[v] = None
            tmp_sz[v] = None
            tmp_dep[v] = None
            tmp_centroid_child_group[v] = []
            tmp_centroid_child_parent[v] = None
        
        return res
    
    def centroid_decomp():
        res = 0
        root_deq = deque([0])
        while root_deq:
            root = root_deq.popleft()
            centroid = find_centroid(root)
            res += calc_sub(centroid)

            #print(centroid,calc_sub(centroid))

            centroid_done[centroid] = True
            for centroid_child in edge[centroid]:
                if not centroid_done[centroid_child]:
                    root_deq.append(centroid_child)
        
        return res
    
    return centroid_decomp() + sum(S)


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

S = [int(c) for c in input()]

print(solve(N,edge,S))
0