def test_size(self) -> None: dsu = DSU(5) dsu.merge(0, 1) assert dsu.size(0) == 2 dsu.merge(0, 2) assert dsu.size(0) == 3 # Test assertion of invalid indices for i in (-1, 5): with pytest.raises(AssertionError): dsu.size(i)
def test_initial_status(self) -> None: dsu = DSU(5) for i, j in combinations(range(5), 2): assert not dsu.same(i, j) for index in range(5): assert dsu.size(index) == 1 assert dsu.leader(index) == index assert dsu.groups() == [[0], [1], [2], [3], [4]]
from atcoder.dsu import DSU N = int(input()) edges = [] for _ in range(N - 1): u, v, w = (int(x) for x in input().split()) edges.append((w, u - 1, v - 1)) edges.sort() dsu = DSU(N) ans = 0 for w, u, v in edges: ans += w * dsu.size(u) * dsu.size(v) dsu.merge(u, v) print(ans)