Pythonで蟻本2-4 - データ構造（ヒープ・二分探索木・Union-Find木）

• プライオリティキュー
  
  • ①数の追加，②最小値を取り出して削除，ができるデータ構造．
  • （二分）ヒープと呼ばれるデータ構造を使うと，要素数$n$の場合に計算量$O(\log n)$で実現できる．
• 二分探索木
  
  • ①数値の追加，②数値が含まれるか調べる，③数値の削除，が効率的に行えるデータ構造．
  • 要素数$n$の場合に計算量$O(\log n)$で実現できる．
• Union-Find木
  
  • ①2要素が同じグループに属するか調べる，②2要素のグループを併合する，が効率的にできるデータ構造．分割はできない．
  • 計算量は$O(\alpha(n))$．

背景：蟻本をやります - 競プロはじめました
テキスト：

プログラミングコンテストチャレンジブック [第2版]　～問題解決のアルゴリズム活用力とコーディングテクニックを鍛える～

【2021.04.22-24】作成期間

プライオリティキュー（順位キュー）

Expedition

【入力】
4
25
10
10 14 20 21
10 5 2 4

【出力】
2

【コード】

import heapq

N = int(input())
L = int(input())
P = int(input())
A = list(map(int, input().split()))
B = list(map(int, input().split()))

# ガソリンスタンドごとに止まって考える．目的地をダミーで追加．
A.append(L)
B.append(0)

que = []
ans, pos, tank = 0, 0, P

for i in range(N):
  d = A[i] - pos
  while (tank < d):
    if len(que) == 0:
      exit(print('-1'))
    tank += -1 * heapq.heappop(que)
    ans += 1
    
  tank -= d
  pos = A[i]
  heapq.heappush(que, -1 * B[i])

print(ans)

【ポイント】
プライオリティキュー（順位キュー）はheapqで使えるようです．しかし，

heapq --- ヒープキューアルゴリズム

(b) われわれの pop メソッドは最大の要素ではなく最小の要素 (教科書では "min heap:最小ヒープ" と呼ばれています; 教科書では並べ替えをインプレースで行うのに適した "max heap:最大ヒープ" が一般的です)。

heapq --- ヒープキューアルゴリズム — Python 3.9.4 ドキュメント

とのことなので，最大値を取り出すためには工夫が必要です．
次の記事のように，追加(push)・取り出し(pop)時に-1を掛けることで処理しました．Pythonで理解する蟻本「2-4 プライオリティキューとヒープ」(p.71) - クルトンのプログラミング教室

Fence Repair

貪欲法で解いた問題をプライオリティキューで解きます．
Pythonで蟻本2-2 - 貪欲法 - 競プロはじめました
【入力】
3
8 5 8

【出力】
34

【コード】

import heapq

n = int(input())
que = list(map(int, input().split()))
heapq.heapify(que)

ans = 0
while len(que) > 1:
  l1 = heapq.heappop(que)
  l2 = heapq.heappop(que)
  
  ans += l1 + l2
  heapq.heappush(que, l1 + l2)
 
print(ans)

【ポイント】

• 今回は，小さい順に取り出すので$\times (-1)$は不要です．

二分探索木

問題がないので略．

Union-Find木

食物連鎖

【入力】
100
7
1 101 1
2 1 2
2 2 3
2 3 3
1 1 3
2 3 1
1 5 5

【出力】
3 ([1, 4, 5])

【コード】

N = int(input())
K = int(input())
info = [list(map(int, input().split())) for _ in range(K)]

# UnionFind
class UnionFind:
  # n要素で初期化
  def __init__(self, n):
    self.par = [i for i in range(n)]
    self.rank = [0]*n
    
  # xの根
  def find(self, x):
    if self.par[x] == x:
      return x
    else:
      self.par[x] = self.find(self.par[x])
      return self.par[x]
 
  # xとyの集合を併合(高い木の根に低い木の根を付ける)
  def unite(self, x, y):
    x = self.find(x)
    y = self.find(y)
    if x == y:
      return
    if self.rank[x] < self.rank[y]:
      self.par[x] = y
    else:
      self.par[y] = x
      if self.rank[x] == self.rank[y]:
        self.rank[x] += 1

  # 同じ集合か(根が同じか)判定
  def same(self, x, y):
    return self.find(x) == self.find(y)

# 処理
uf = UnionFind(N * 3)
ans = 0
ans_list = []  # 追加
for i in range(K):
  t = info[i][0]
  x = info[i][1] - 1
  y = info[i][2] - 1
  
  if x < 0 or N <= x or y < 0 or N <= y:
    # 範囲外の場合
    ans += 1
    ans_list.append(i+1)  # 追加
    # 以降の処理は飛ばす
    continue

  if t == 1:
    # 同じ種類の情報
    # (A,B,Cのどれか不明→全ての可能性を並行処理)
    # x-Aとy-B,Cが同時に起きないか確認
    # A,B,Cを同等に処理しているので，x-B,Cの確認は不要
    # (y, y+N, y+2Nがy-A, y-B, y-Cに対応)
    if uf.same(x, y+N) or uf.same(x, y+2*N):
      ans += 1
      ans_list.append(i+1)  # 追加
    else:
      # 矛盾しないならx-A,B,Cとy-A,B,Cを併合
      uf.unite(x, y)
      uf.unite(x+N, y+N)
      uf.unite(x+2*N, y+2*N)
  else:
    # (A,B,Cのどれか不明→全ての可能性を並行処理)
    # 食べるという情報→x-A, y-Bでなければダメ
    # A,B,Cを同等に処理しているので，
    # (x-B & y-C), (x-C & y-A)の確認不要
    if uf.same(x, y) or uf.same(x, y+2*N):
      ans += 1
      ans_list.append(i+1)  # 追加
    else:
      uf.unite(x, y+N)
      uf.unite(x+N, y+2*N)
      uf.unite(x+2*N, y)

print('{} ({})'.format(ans, ans_list))  

【ポイント】
今回は，コメントとして記入してみました．

• $A,B,C$のどれかという情報はないので，全ての可能性を並行して処理するために3倍の要素を考えています．
  
  • $x,x+N,x+2N$のそれぞれが，$x$が$A$($x$-$A$)だとしたら〜 ，$B$($x$-$B$)だとしたら〜 ，$C$($x$-$C$)だとしたら〜，に対応しています．