結果

問題 No.1843 Tree ANDistance
ユーザー lam6er
提出日時 2025-03-20 20:30:24
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 682 ms / 2,000 ms
コード長 2,143 bytes
コンパイル時間 243 ms
コンパイル使用メモリ 82,540 KB
実行使用メモリ 253,872 KB
最終ジャッジ日時 2025-03-20 20:31:00
合計ジャッジ時間 15,821 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other AC * 38
権限があれば一括ダウンロードができます

ソースコード

diff # 

MOD = 10**9 + 7

def main():
    import sys
    input = sys.stdin.read().split()
    ptr = 0
    N = int(input[ptr])
    ptr += 1
    
    bit_edges = [[] for _ in range(31)]
    for _ in range(N - 1):
        a = int(input[ptr]) - 1
        b = int(input[ptr + 1]) - 1
        c = int(input[ptr + 2])
        ptr += 3
        for k in range(31):
            if (c >> k) & 1:
                bit_edges[k].append((a, b))
    
    ans = 0
    for k in range(31):
        parent = list(range(N))
        size = [1] * N
        
        for a, b in bit_edges[k]:
            # Find roots with path compression
            root_a = a
            while parent[root_a] != root_a:
                parent[root_a] = parent[parent[root_a]]  # Path compression during find
                root_a = parent[root_a]
            root_b = b
            while parent[root_b] != root_b:
                parent[root_b] = parent[parent[root_b]]  # Path compression during find
                root_b = parent[root_b]
            
            if root_a != root_b:
                if size[root_a] < size[root_b]:
                    parent[root_a] = root_b
                    size[root_b] += size[root_a]
                else:
                    parent[root_b] = root_a
                    size[root_a] += size[root_b]
        
        # After all unions, find all roots and their sizes
        # We need to perform path compression again to ensure all nodes point directly to their root
        roots = {}
        for i in range(N):
            x = i
            # Find root with path compression
            while parent[x] != x:
                parent[x] = parent[parent[x]]  # Compress path
                x = parent[x]
            root = x
            # Now, update the parent of i to root (compression)
            if parent[i] != root:
                parent[i] = root
            if root not in roots:
                roots[root] = size[root]
        
        total = 0
        for s in roots.values():
            total += s * (s - 1) // 2
        
        ans = (ans + (total << k)) % MOD
    
    print(ans % MOD)

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