結果

問題 No.399 動的な領主
ユーザー kept1994kept1994
提出日時 2023-03-20 20:36:02
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 869 ms / 2,000 ms
コード長 7,722 bytes
コンパイル時間 257 ms
コンパイル使用メモリ 82,264 KB
実行使用メモリ 188,032 KB
最終ジャッジ日時 2024-09-18 14:10:35
合計ジャッジ時間 10,171 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge2
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 63 ms
67,840 KB
testcase_01 AC 58 ms
67,328 KB
testcase_02 AC 63 ms
69,504 KB
testcase_03 AC 63 ms
69,760 KB
testcase_04 AC 122 ms
80,000 KB
testcase_05 AC 223 ms
83,584 KB
testcase_06 AC 869 ms
111,104 KB
testcase_07 AC 802 ms
111,744 KB
testcase_08 AC 796 ms
114,304 KB
testcase_09 AC 806 ms
114,304 KB
testcase_10 AC 119 ms
80,000 KB
testcase_11 AC 173 ms
83,328 KB
testcase_12 AC 657 ms
122,968 KB
testcase_13 AC 647 ms
122,112 KB
testcase_14 AC 440 ms
187,904 KB
testcase_15 AC 461 ms
188,032 KB
testcase_16 AC 529 ms
157,056 KB
testcase_17 AC 868 ms
113,024 KB
testcase_18 AC 808 ms
114,176 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

#!/usr/bin/env python3
from typing import List, Tuple
from collections import deque
import sys
sys.setrecursionlimit(10 ** 9)

class BIT:
    def __init__(self, N):
        self.N = N
        self.A = [0] * (N + 1)

    def build(self, A):
        """build BIT with given list"""
        for i, a in enumerate(A):
            self.A[i + 1] = a
        for i in range(1, self.N):
            if i + (i & -i) > self.N:
                continue
            self.A[i + (i & -i)] += self.A[i]

    def add(self, i, x):
        """add x to i-th element (0-indexed)"""
        # assert 0 <= i < self.N
        i += 1
        while i <= self.N:
            self.A[i] += x
            i += i & -i

    def sum(self, i):
        """return sum(A[:i])"""
        assert 0 <= i <= self.N
        s = 0
        while i > 0:
            s += self.A[i]
            i -= i & -i
        return s

    def range_sum(self, l, r):
        """return sum(A[l:r])"""
        return self.sum(r) - self.sum(l)

def main():
    class HLD:
        def __init__(self, N, G, root: int = 0):
            self.N = N                  # ノード数
            self.G = G                  # グラフ(隣接リスト)
            self.root = root            # 根
            self.depth_of_node = [0] * self.N                 # 各ノードの根からの距離(深さ)
            self.next_heavy_nodes = [-1] * self.N             # next_heavy_nodes[i] := heavy-pathにおいてノードiの次ノード番号 (0-indexed)
            self.parent_nodes = [None] * self.N               # parent_nodes[i] := ノードiの親ノード番号 (0-indexed)
            self.partial_tree_size = [None] * self.N          # partial_tree_size[i] := ノードi以下の部分木のサイズ

            self.heavy_paths = []                            # 全heavy_pathのリスト
            self.depth_of_heavy_path = []                    # depth_of_heavy_path[i] := i番目のheavy_pathの深さ(幾つのheavy_pathを乗り継いで到達するか)
            # Depth 0: [0] -------------------- [2] - [5] - [6] - [8]
            #           |                        |     |
            # Depth 1: [1] - [3] - [9] - [10]   [4]   [7] 
            #                       |
            # Depth 2:             [11]
            self.to_heavy_path_index = [0] * N               # ノードiはどのheavy_pathに所属するか(heavy_pathsのインデックス)
            self.indices_in_heavy_path = [0] * N             # ノードiがその属するheavy-pathの何番目にあたるのか。
            self.heads = [None] * self.N                     # ノードiが属するheavy-pathの先頭のノード番号
            self.ord = [None] * self.N                       # heavy-pathに割り当てた連続する番号
            # 上の例において
            # [0] -------------------- [5] - [6] - [7] - [8]
            #  |                        |     |
            # [1] - [2] - [3] - [4]    [10]  [11] 
            #              |
            #             [9]

        def path_query_range(self, a: int, b: int) -> List[Tuple[int, int]]:
            """
            ノードaとノードbの間に含まれるordの範囲[l, r)のリストを返す。
            """
            ret = []
            while True:
                if self.ord[a] > self.ord[b]:
                    a, b = b, a
                if self.heads[a] == self.heads[b]:
                    ret.append((self.ord[a], self.ord[b] + 1))
                    return ret
                ret.append((self.ord[self.heads[b]], self.ord[b] + 1))
                b = self.parent_nodes[self.heads[b]]

        # def subtree_query_range(self, a: int) -> Tuple[int, int]:
        #     """
        #     ノードaの部分木に含まれるordの範囲[l, r)のリストを返す。
            
        #     return [l, r) range that cover vertices of subtree v"""
        #     return (self.ord[a], self.ord[a] + self.partial_tree_size[a])

        def lca(self, u, v):
            while True:
                if self.ord[u] > self.ord[v]:
                    u, v = v, u
                if self.heads[u] == self.heads[v]:
                    return u
                v = self.parent_nodes[self.heads[v]]

        def _dfs(self, cur, depth):
            self.depth_of_node[cur] = depth
            sub_node_count = 1              # そのノード(を含む)配下の子ノード数
            heavy_path = None               # そのノードからのheavy-path
            max_heavy_path_length = 0       # 最大のheavy-path長
            for next in G[cur]:
                if self.parent_nodes[cur] == next:
                    continue
                self.parent_nodes[next] = cur
                sub_node_total_count = self._dfs(next, depth + 1)
                # より長いパスが見つかったなら
                # heavy-pathとしてその子ノードの番号を仮置きする。
                if max_heavy_path_length < sub_node_total_count:
                    heavy_path = next
                    max_heavy_path_length = sub_node_total_count
                sub_node_count += sub_node_total_count
            self.partial_tree_size[cur] = sub_node_count
            self.next_heavy_nodes[cur] = heavy_path
            return sub_node_count

        def build(self):
            """
            グラフGをheavy-pathに沿って縮約
            """
            self._dfs(self.root, 0)
            q = deque([(0, 0)])
            order = 0
            while q:
                cur, depth_of_cur_heaby_path = q.popleft()
                head_of_heavy_path = cur
                cur_heavy_path = []
                k = len(self.heavy_paths)
                while cur is not None:
                    self.ord[cur] = order
                    self.heads[cur] = head_of_heavy_path
                    self.indices_in_heavy_path[cur] = len(cur_heavy_path)
                    cur_heavy_path.append(cur)
                    self.to_heavy_path_index[cur] = k
                    next_heavy_node = self.next_heavy_nodes[cur]

                    # heavy-path以外をキューに入れておき、heavy-path上の次のノードの処理へ。
                    for next in G[cur]:
                        if self.parent_nodes[cur] == next or next_heavy_node == next:
                            continue
                        q.append((next, depth_of_cur_heaby_path + 1))
                    cur = next_heavy_node
                    order += 1
                    
                self.heavy_paths.append(cur_heavy_path)
                self.depth_of_heavy_path.append(depth_of_cur_heaby_path)

    # N = 8
    # M = 7
    # A = [0, 0, 0, 1, 1, 5, 5]
    # B = [1, 2, 3, 4, 5, 6, 7]
    # N = 6
    # M = 5
    # A = [0, 1, 1, 0, 4]
    # B = [1, 2, 3, 4, 5]
    N = int(input())
    G = [[] for _ in range(N)]
    for i in range(N - 1):
        u, v = map(int, input().split())
        u -= 1
        v -= 1
        G[u].append(v)
        G[v].append(u)

    hl = HLD(N, G)
    hl.build()

    # print(hl.parent_nodes)
    # print(hl.partial_tree_size)
    # print(hl.heads)
    # print(hl.heavy_paths)
    # print(hl.depth_of_heavy_path)
    # print(hl.to_heavy_path_index)
    # print(hl.indices_in_heavy_path)
    # print(hl.depth_of_node)
    # print(hl.ord)

    bit = BIT(N)
    Q = int(input())
    for _ in range(Q):
        a, b = map(int, input().split())
        a -= 1
        b -= 1
        for l, r in hl.path_query_range(a, b):
            bit.add(l, 1)
            bit.add(r, -1)
    ans = 0
    for i in range(N):
        cnt = bit.sum(i + 1)
        ans += (1 + cnt) * cnt // 2
    print(ans)

if __name__ == '__main__':
    main()
0