結果

問題 No.1054 Union add query
コンテスト
ユーザー hogeandfuga
提出日時 2021-04-25 14:34:03
言語 Python3
(3.13.1 + numpy 2.2.1 + scipy 1.14.1)
結果
TLE  
実行時間 -
コード長 2,099 bytes
コンパイル時間 77 ms
コンパイル使用メモリ 12,800 KB
実行使用メモリ 181,740 KB
最終ジャッジ日時 2024-07-04 09:27:33
合計ジャッジ時間 13,219 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 7 TLE * 1
権限があれば一括ダウンロードができます

ソースコード

diff # 

from typing import Any


class WeightedUnionFind():

    T = Any

    def __init__(self, n: int) -> None:
        self.num = n
        self.rs = [1]*n
        self.ps = list(range(n))
        self.ws = [0]*n
        self.gs = [{i} for i in range(n)]

    def find(self, x: int) -> int:
        if x == self.ps[x]:
            return x
        else:
            t = self.find(self.ps[x])
            self.ws[x] += self.ws[self.ps[x]]
            self.ps[x] = t
            return self.ps[x]

    def weight(self, x: int) -> T:
        self.find(x)
        return self.ws[x]

    def same(self, x: int, y: int) -> bool:
        return self.find(x) == self.find(y)

    def unite(self, x: int, y: int, w: T) -> None:
        w += self.weight(x)
        w -= self.weight(y)
        x, y = self.find(x), self.find(y)
        if x == y:
            return
        if self.rs[x] < self.rs[y]:
            x, y = y, x
            w = -w
        self.rs[x] += self.rs[y]
        self.ps[y] = x
        self.ws[y] = w
        self.gs[x].update(self.gs[y])
        self.num -= 1

    def diff(self, x: int, y: int) -> T:
        return self.weight(y) - self.weight(x)

    def size(self, x: int) -> int:
        return self.rs[self.find(x)]

    def count(self) -> int:
        return self.num

    def group(self, x: int) -> set:
        return self.gs[self.find(x)]

import sys
input = sys.stdin.buffer.readline

def main() -> None:
    N, Q = map(int, input().split())
    uf = WeightedUnionFind(N)
    vs = [0]*N

    def add(a: int, b: int) -> None:
        vs[uf.find(a)] += b

    def query(a: int) -> int:
        r = uf.find(a)
        return uf.diff(a, r) + vs[r]

    def unite(a: int, b: int) -> None:
        a = uf.find(a)
        b = uf.find(b)
        if a == b:
            return
        uf.unite(a, b, vs[a] - vs[b])

    for _ in range(Q):
        t, a, b = map(int, input().split())
        a -= 1
        if t == 1:
            b -= 1
            unite(a, b)

        if t == 2:
            add(a, b)
        if t == 3:
            print(query(a))


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