DSL_1の解答例(Python)

素因数分解・累乗計算・最小公倍数・オイラーのφ関数・拡張ユークリッド互除法に関する問題です．

1. DSL_1_A: Disjoint Set: Union Find Tree

# DSL_1_A 単純なコード
# 本来はUnionFindを使う
# sameやunite で find を2回呼び出すので効率が悪い
# コード全体での計算量は O(2*q*n) 本当は 0 n-1 n-2 を何度も呼び出すと2*10**9程度かかってしまう

n,q = map(int,input().split())
set_list = [set([i]) for i in range(n)]

# xを含むセットの場所を返す : O(n)
def find_set(x):
    for i,S in enumerate(set_list):
        if x in S:
            return i
    
# unite : 2つの集合の和集合 : O(2*n)
def unite_set(x,y):
    idx_x = find_set(x)
    idx_y = find_set(y)
    # 違う集合に含まれる場合のみ操作を行えばよい
    if idx_y != idx_x:
        # インデックスの位置がずれないように大きいほうから取り出す
        set_x = set_list.pop(max(idx_x,idx_y))
        set_y = set_list.pop(min(idx_x,idx_y))
        new_set = set_x | set_y
        set_list.append(new_set)
        
# same : 同じ集合に含まれるかどうか O(2*n)
def is_same(x,y):
    idx_x = find_set(x)
    idx_y = find_set(y)
    if idx_x == idx_y:
        print(1)
    else:
        print(0)
    
for _ in range(q):
    com,x,y = map(int,input().split())
    # unite : com == 0
    if com == 0:
        unite_set(x,y)
    # same : com == 1
    elif com == 1:
        is_same(x,y)

# DSL_1_A 単純なUnionFind
class UnionFind:
    def __init__(self,n):
        self.parent_dict = {i:None for i in range(n)}
        
    # 根を探す
    def root(self,x):
        # xが根の場合
        if self.parent_dict[x] == None:
            return x
        else:
            self.parent_dict[x] = self.root(self.parent_dict[x]) # 経路圧縮
            return self.parent_dict[x]
       
    # 同じ集合に属するかを調べる = 根が同じかを調べる
    def isSame(self,x,y):
        return self.root(x) == self.root(y)
        
    def unite(self,x,y):
        if not self.isSame(x,y):
            self.parent_dict[self.root(y)] = self.root(x)
    
n,q = map(int,input().split())
UF = UnionFind(n)
for _ in range(q):
    com,x,y = map(int,input().split())
    # unite : com == 0
    if com == 0:
        UF.unite(x,y)
    # same : com == 1
    elif com == 1:
        print(int(UF.isSame(x,y)))

# 木を平衡に保つためにランクを導入している
class UnionFind:
    def __init__(self, n):
        # 親を格納するリストと各要素のランクを格納するリストを初期化する
        self.parent = [i for i in range(n)]  # 最初は自分自身が親
        self.rank = [0] * n  # ランクは最初すべて0
    
    def find(self, x):
        # 要素xの親を見つける
        if self.parent[x] != x:
            # 要素の親が自分自身でない場合、再帰的に親を探す
            self.parent[x] = self.find(self.parent[x])  # 親を再帰的に探して更新
        return self.parent[x]
    
    def unite(self, x, y):
        # xとyのルートを見つける
        root_x = self.find(x)
        root_y = self.find(y)
        
        if root_x == root_y:
            # xとyが同じグループに属している場合は何もしない
            return None
        
        # ランクを考慮して結合する(小さい木を大きい木に結合する)
        if self.rank[root_x] < self.rank[root_y]:
            # root_xのランクがroot_yのランクより小さい場合、root_xをroot_yの子にする
            self.parent[root_x] = root_y
        elif self.rank[root_x] > self.rank[root_y]:
            # root_xのランクがroot_yのランクより大きい場合、root_yをroot_xの子にする
            self.parent[root_y] = root_x
        else:
            # 両方のランクが同じ場合はどちらを親にしても良いが、ここではroot_yをroot_xの子にする
            self.parent[root_y] = root_x
            self.rank[root_x] += 1  # root_xのランクを増やす
            
    def isSame(self,x,y):
        return self.find(x) == self.find(y)
            
n,q = map(int,input().split())
UF = UnionFind(n)
for _ in range(q):
    com,x,y = map(int,input().split())
    # unite : com == 0
    if com == 0:
        UF.unite(x,y)
    # same : com == 1
    elif com == 1:
        print(int(UF.isSame(x,y)))

2. DSL_1_B: Weighted Union Find Trees

# DSL_1_B
class UnionFind:
    def __init__(self, n):
        self.parent_dict = {i: None for i in range(n)}
        self.diff_dict = {i: 0 for i in range(n)}  # 親との差分の辞書 0で初期化しておく(根の重みは常に0)
        
    # 根を探す
    def root(self, x):
        # xが根の場合
        if self.parent_dict[x] is None:
            return x
        else:
            root = self.root(self.parent_dict[x])  # 経路圧縮
            self.diff_dict[x] += self.diff_dict[self.parent_dict[x]]  # 根を基準に差分を累積していく
            self.parent_dict[x] = root  # 経路圧縮
            return root
       
    # 同じ集合に属するなら差が求まる
    def get_diff(self, x, y):
        if self.root(x) != self.root(y):
            return "?"
        return self.diff_dict[y] - self.diff_dict[x]

    def relate(self, x, y, z):
        root_x = self.root(x)
        root_y = self.root(y)

        if root_x == root_y:
            return

        self.parent_dict[root_y] = root_x
        self.diff_dict[root_y] = self.diff_dict[x] + z - self.diff_dict[y]  # yの親との差分を更新

    # 入力処理
n, q = map(int, input().split())
uf = UnionFind(n)

# クエリの処理
for _ in range(q):
    query = input().split()
    if query[0] == '0':
        x, y, z = map(int, query[1:])
        uf.relate(x, y, z)
    else:
        x, y = map(int, query[1:])
        result = uf.get_diff(x, y)
        print(result)