結果

問題 No.1054 Union add query
ユーザー hogeandfugahogeandfuga
提出日時 2021-04-25 15:07:58
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 797 ms / 2,000 ms
コード長 2,101 bytes
コンパイル時間 278 ms
コンパイル使用メモリ 82,304 KB
実行使用メモリ 235,368 KB
最終ジャッジ日時 2024-07-04 09:27:54
合計ジャッジ時間 6,082 ms
ジャッジサーバーID
(参考情報)
judge2 / judge3
このコードへのチャレンジ
(要ログイン)

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 82 ms
67,072 KB
testcase_01 AC 66 ms
67,072 KB
testcase_02 AC 70 ms
67,200 KB
testcase_03 AC 670 ms
124,352 KB
testcase_04 AC 797 ms
235,368 KB
testcase_05 AC 515 ms
102,988 KB
testcase_06 AC 345 ms
154,380 KB
testcase_07 AC 307 ms
153,996 KB
testcase_08 AC 321 ms
153,996 KB
testcase_09 AC 447 ms
231,296 KB
testcase_10 AC 336 ms
229,888 KB
権限があれば一括ダウンロードができます

ソースコード

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(r, a) + 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[b] - vs[a])

    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