import sys import typing import dataclasses import collections _readline = sys.stdin.readline _write = sys.stdout.write _flush = sys.stdout.flush _print = print @dataclasses.dataclass class Node: rng: int edges: typing.Dict[int, int] = dataclasses.field(default_factory=dict) @dataclasses.dataclass class Tree: nodes: typing.Dict[int, Node] = dataclasses.field(default_factory=dict) def print(*args): _print(*args, file=sys.stderr) def fire_node(tree: Tree, node_id: int) -> typing.Dict[int, bool]: node = tree.nodes[node_id] stack = collections.deque([(node, node.rng)]) result = {node_id: True} while stack: node, rng_remaining = stack.pop() for v, l in node.edges.items(): if rng_remaining >= l: result[v] = True next_node = tree.nodes[v] stack.append((next_node, rng_remaining-l)) return result def get_fired_nodes_count(tree: Tree, node_id: int, memo: dict) -> int: queue = collections.deque([node_id]) visited = {} while queue: current_node_id = queue.popleft() if current_node_id not in visited: visited[current_node_id] = True if current_node_id not in memo: memo[current_node_id] = fire_node(tree, current_node_id) fired_by_current_node = memo[current_node_id] for v in fired_by_current_node: if v not in visited: queue.append(v) return len(visited) def solve(tree: Tree) -> typing.List[int]: out = [] memo = {} for i in range(1, len(tree.nodes)+1): out.append(get_fired_nodes_count(tree, i, memo)) return out def main(): def read_int() -> int: return int(_readline().strip()) def read_int_array() -> typing.List[int]: return [int(x.strip()) for x in _readline().split()] def write_int_array(values: typing.List[int]): for v in values: _write(f'{v} ') _flush() tree = Tree() nodes = tree.nodes n = read_int() ranges = read_int_array() for i, r in enumerate(ranges): nodes[i+1] = Node(r) for _ in range(n-1): va, vb, l = read_int_array() nodes[va].edges[vb] = l nodes[vb].edges[va] = l result = solve(tree) write_int_array(result) if __name__ == '__main__': main()
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | import sys import typing import dataclasses import collections _readline = sys.stdin.readline _write = sys.stdout.write _flush = sys.stdout.flush _print = print @dataclasses.dataclass class Node: rng: int edges: typing.Dict[int, int] = dataclasses.field(default_factory=dict) @dataclasses.dataclass class Tree: nodes: typing.Dict[int, Node] = dataclasses.field(default_factory=dict) def print(*args): _print(*args, file=sys.stderr) def fire_node(tree: Tree, node_id: int) -> typing.Dict[int, bool]: node = tree.nodes[node_id] stack = collections.deque([(node, node.rng)]) result = {node_id: True} while stack: node, rng_remaining = stack.pop() for v, l in node.edges.items(): if rng_remaining >= l: result[v] = True next_node = tree.nodes[v] stack.append((next_node, rng_remaining-l)) return result def get_fired_nodes_count(tree: Tree, node_id: int, memo: dict) -> int: queue = collections.deque([node_id]) visited = {} while queue: current_node_id = queue.popleft() if current_node_id not in visited: visited[current_node_id] = True if current_node_id not in memo: memo[current_node_id] = fire_node(tree, current_node_id) fired_by_current_node = memo[current_node_id] for v in fired_by_current_node: if v not in visited: queue.append(v) return len(visited) def solve(tree: Tree) -> typing.List[int]: out = [] memo = {} for i in range(1, len(tree.nodes)+1): out.append(get_fired_nodes_count(tree, i, memo)) return out def main(): def read_int() -> int: return int(_readline().strip()) def read_int_array() -> typing.List[int]: return [int(x.strip()) for x in _readline().split()] def write_int_array(values: typing.List[int]): for v in values: _write(f'{v} ') _flush() tree = Tree() nodes = tree.nodes n = read_int() ranges = read_int_array() for i, r in enumerate(ranges): nodes[i+1] = Node(r) for _ in range(n-1): va, vb, l = read_int_array() nodes[va].edges[vb] = l nodes[vb].edges[va] = l result = solve(tree) write_int_array(result) if __name__ == '__main__': main() |