Problem p02710 (python)

Fixed submission:
import sys
import itertools

sys.setrecursionlimit(3 * 10**5)


class Node:
  def __init__(self, color):
    self.color = int(color) - 1
    self.children = []

  def add_child(self, node):
    self.children.append(node)

  @staticmethod
  def prepare_nodes(n):
    nodes = tuple(map(Node, input().split()))
    for _ in range(n - 1):
      a, b = map(lambda x: int(x) - 1, input().split())
      nodes[a].add_child(nodes[b])
      nodes[b].add_child(nodes[a])

    return nodes


def main():
  n = int(input())
  nodes = Node.prepare_nodes(n)

  dp = [0 for _ in range(n)]
  ans_list = [n*(n + 1) // 2 for _ in range(n)]

  def update_ans(s, color):
    ans_list[color] -= (s + dp[color]) * (s + dp[color] + 1) // 2

  def dfs(s, node, parent):
    color = node.color
    pre_s = s + dp[color]
    for child in node.children:
      if child == parent: continue

      dp[color] = -s
      s = dfs(s, child, node)
      update_ans(s, color)

    s += 1
    dp[color] = pre_s - s

    return s

  s = dfs(0, nodes[0], Node(-1))

  for i in range(n):
    update_ans(s, i)
    print(ans_list[i])


main()

Buggy submission (verdict: Runtime Error):
import sys
import itertools

sys.setrecursionlimit(3 * 10**5)


class Node:
  def __init__(self, tpl):
    self.id = tpl[0] + 1
    self.c = int(tpl[1])
    self.children = []

  def add_child(self, node):
    self.children.append(node)

  @staticmethod
  def prepare_nodes(n):
    nodes = tuple(map(Node, enumerate(input().split())))
    for _ in range(n - 1):
      a, b = map(lambda x: int(x) - 1, input().split())
      nodes[a].add_child(nodes[b])
      nodes[b].add_child(nodes[a])

    return nodes


def main():
  n = int(input())
  nodes = Node.prepare_nodes(n)

  dp = [0 for _ in range(n)]
  ans_list = [n*(n + 1) // 2 for _ in range(n)]

  def update_ans(s, color):
    ans_list[color] -= (s + dp[color]) * (s + dp[color] + 1) // 2

  def dfs(s, node, parent):
    color = node.c
    pre_s = s + dp[color]
    for child in node.children:
      if child == parent: continue

      dp[color] = -s
      s = dfs(s, child, node)
      update_ans(s, color)

    s += 1
    dp[color] = pre_s - s

    return s

  s = dfs(0, nodes[0], Node((-1, -1)))

  for i in range(n):
    update_ans(s, i)

  print(*ans_list, sep='\n')


main()

Buggy submission (verdict: Runtime Error):
import sys

sys.setrecursionlimit(4000)


class Node:
  def __init__(self, tpl):
    self.id = tpl[0] + 1
    self.c = int(tpl[1])
    self.relations = []

  def add_relation(self, node):
    self.relations.append(node)


def prepare_nodes(n):
  node_dict = {
    node.id: node
    for node in map(Node, enumerate(input().split()))
  }

  for _ in range(n - 1):
    a, b = map(int, input().split())
    node_dict[a].add_relation(node_dict[b])
    node_dict[b].add_relation(node_dict[a])

  return node_dict


def get_connected_node_ids(k, node, prev_id = None):
  other_color_nodes = filter(
    lambda node: node.c != k and node.id != prev_id,
    node.relations
  )

  result = {node.id}
  for next_node in other_color_nodes:
    result |= get_connected_node_ids(k, next_node, node.id)
  return result


def main():
  n = int(input())
  node_dict = prepare_nodes(n)
  comb_num = int(n * (n - 1) / 2) + n

  for k in range(1, n + 1):
    other_color_nodes = tuple(filter(
      lambda node: node.c != k, node_dict.values()
    ))
    num = len(other_color_nodes)
    if num == n:
      print(0)
      continue

    done_ids = set()
    for node in other_color_nodes:
      if node.id in done_ids:
        continue

      node_ids = get_connected_node_ids(k, node)
      node_ids_len = len(node_ids)

      num += node_ids_len * (node_ids_len - 1) / 2
      done_ids |= node_ids

    print(comb_num - int(num))


main()

Buggy submission (verdict: Runtime Error):
import sys

sys.setrecursionlimit(4000)


class Node:
  def __init__(self, tpl):
    self.id = tpl[0] + 1
    self.c = int(tpl[1])
    self.relations = []

  def add_relation(self, node):
    self.relations.append(node)


def prepare_nodes(n):
  node_dict = {
    node.id: node
    for node in map(Node, enumerate(input().split()))
  }

  for _ in range(n - 1):
    a, b = map(int, input().split())
    node_dict[a].add_relation(node_dict[b])
    node_dict[b].add_relation(node_dict[a])

  return node_dict


def cal(k, node, prev_id = None):
  other_color_nodes = list(filter(
    lambda node: node.c != k and node.id != prev_id,
    node.relations
  ))
  num = len(other_color_nodes)

  for next_node in other_color_nodes:
    num += cal(k, next_node, node.id)
  return num


def main():
  n = int(input())
  node_dict = prepare_nodes(n)
  comb_num = int(n * (n - 1) / 2) + n

  for k in range(1, n + 1):
    other_color_nodes = tuple(filter(
      lambda node: node.c != k, node_dict.values()
    ))
    num = len(other_color_nodes)
    if num == n:
      print(0)
      continue

    num *= 2
    for node in other_color_nodes:
      num += cal(k, node)

    print(comb_num - int(num / 2))


main()

Buggy submission (verdict: Runtime Error):
import sys
from collections import defaultdict

sys.setrecursionlimit(4000)


class Node:
  def __init__(self, tpl):
    self.id = tpl[0] + 1
    self.c = int(tpl[1])
    self.relations = {}

  def add_relation(self, node):
    self.relations[node.id] = node

  def print_node_with_relations(self, nest = 0, exclude_id = None):
    print("  " * nest, self.get_print_str())
    for node in self.relations.values():
      if node.id == exclude_id: continue
      node.print_node_with_relations(nest + 1, self.id)

  def print_node(self):
    print(self.get_print_str())

  def get_print_str(self):
    return f"{self.id}({self.c})"


def prepare_nodes(n):
  node_dict = {
    node.id: node
    for node in map(Node, enumerate(input().split()))
  }

  for _ in range(n - 1):
    a, b = sorted(map(int, input().split()))
    node_dict[a].add_relation(node_dict[b])
    node_dict[b].add_relation(node_dict[a])

  return node_dict


def cal(k, node, prev_id = None):
  other_color_nodes = list(filter(
    lambda node: node.c != k and node.id != prev_id,
    node.relations.values()
  ))
  num = len(other_color_nodes)
  for next_node in other_color_nodes:
    num += cal(k, next_node, node.id)
  return num


def main():
  n = int(input())
  node_dict = prepare_nodes(n)

  for k in range(1, n + 1):
    other_color_nodes = tuple(filter(
      lambda node: node.c != k, node_dict.values()
    ))
    num = len(other_color_nodes) * 2
    if num == n * 2:
      print(0)
      continue

    for node in other_color_nodes:
      num += cal(k, node)

    print(int(n * (n - 1) / 2) + n - int(num / 2))


main()