【Python3】【殿堂入りキャンペーン】村人の友好関係 (paizaランク S 相当)：アルゴリズム付きの解説

この問題では、村人の同好会グループの管理と人気度の計算を効率的に行う必要があります。Union-Findデータ構造を用いることで、友好関係を管理し、入退会ログに基づいて人気度を計算することができます。

アルゴリズム設計

概要:

• 村人間の友好度データを受け取り、Union-Findを用いて管理します。
• 各村人が同好会に入退会するたびに、同好会の人気度を計算します。

手順

1. Union-Findクラスの定義:

   • findメソッドでルートを探し、パス圧縮を行う。
   • unionメソッドで2つの集合を結合する。

2. データの読み込み:

   • 村人の数 N、友好関係の数 M、ログの数 Q を読み取る。
   • 友好関係のデータをリスト all_edges に格納し、友好度の降順でソートする。

3. 最小全域木（MST）の構築:

   • Union-Findを用いて友好関係をMSTに追加する。

4. 同好会の入退会の管理:

   • 入退会ログを処理し、同好会に所属する村人の状態を管理するリスト is_member を更新する。
   • 入退会ごとにMSTを用いて人気度を計算する。

完成したコード

# utf-8
class UnionFind:
    def __init__(self, size):
        self.par = [-1 for _ in range(size)]
    
    def find(self, x):
        if self.par[x] < 0:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]
    
    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)
        
        if x == y: return False
        if self.par[x] <= self.par[y]:
            self.par[x] += self.par[y]
            self.par[y] = x
        else:
            self.par[y] += self.par[x]
            self.par[x] = y
        
        return True
        

if __name__ == "__main__":
    import sys
    input = sys.stdin.read
    data = input().split()
    
    index = 0
    N = int(data[index])
    M = int(data[index + 1])
    Q = int(data[index + 2])
    index += 3
    
    all_edges = []
    for _ in range(M):
        a = int(data[index]) - 1
        b = int(data[index + 1]) - 1
        f = int(data[index + 2])
        index += 3
        all_edges.append((f, a, b))
    
    all_edges.sort(reverse=True)
    
    uf = UnionFind(N)
    mst_edges = []
    for f, a, b in all_edges:
        if uf.union(a, b):
            mst_edges.append((a, b, f))
    
    is_member = [False] * N
    answers = []
    
    for _ in range(Q):
        op = data[index]
        member = int(data[index + 1]) - 1
        index += 2
        
        if op == "+":
            is_member[member] = True
        else:
            is_member[member] = False
        
        answer = 0
        for a, b, f in mst_edges:
            if is_member[a] != is_member[b]:
                answer = max(answer, f)
        
        answers.append(answer)
    
    sys.stdout.write("\n".join(map(str, answers)) + "\n")

説明

• Union-Findクラス:

  • __init__メソッドで初期化し、findメソッドでルートを見つけ、unionメソッドで集合を結合します。

• データの読み込みとMSTの構築:

  • 入力データを読み込み、友好関係を降順にソートします。
  • Union-Findを用いてMSTを構築します。

• 同好会の入退会の管理と人気度の計算:

  • is_memberリストで各村人の同好会所属状態を管理します。
  • 入退会ごとにMSTを参照し、同好会内外の最大友好度を計算します。

このアルゴリズムは、各ログエントリに対して効率的に同好会の人気度を計算することができます。