問題

エネルギー最適化問題の例題として、以下のような問題を考えてみましょう。

🔹2つの発電所があり、それぞれ最大で100MWと50MWの電力を供給できる。
🔹需要は150MWである。
🔹各発電所の発電コストは、1MWあたりそれぞれ10ドルと20ドルである。
🔹目的は、発電コストを最小化することである。

解き方・ソースコード

PythonのPuLPライブラリを使用してこの問題を解いてみます。

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

import pulp

# 問題を定義
problem = pulp.LpProblem('Optimization Problem', pulp.LpMinimize)

# 変数を定義
p1 = pulp.LpVariable('p1', lowBound=0, upBound=100, cat='Continuous')
p2 = pulp.LpVariable('p2', lowBound=0, upBound=50, cat='Continuous')

# 目的関数を定義
problem += 10*p1 + 20*p2, 'Total Cost'

# 制約条件を定義
problem += p1 + p2 >= 150, 'Demand'
problem += p1 <= 100, 'Capacity of Plant 1'
problem += p2 <= 50, 'Capacity of Plant 2'

# 問題を解く
problem.solve()

# 結果を出力
print(f'Optimal solution found with total cost of {problem.objective.value():.2f}')
print(f'Plant 1 generates {p1.value()} MW')
print(f'Plant 2 generates {p2.value()} MW')

まずpulpで問題を定義します。

その後、発電所1と発電所2の発電量を表す変数p1とp2を定義します。

次に、目的関数を定義します。

この問題では、発電コストを最小化する必要があるため、変数p1とp2に対して、それぞれの発電コストを乗じて、総和をとります。

この問題では、需要を満たす必要があるため、発電所1と発電所2の発電量の合計が150MW以上でなければなりません。

また、各発電所の発電量は、それぞれ最大で100MWと50MWでなければなりません。

解には、発電所1と発電所2の発電量が含まれます。

[実行結果]

Optimal solution found with total cost of 2000.00

Plant 1 generates 100.0 MW

Plant 2 generates 50.0 MW

最適な発電コストは2000.00であり、その時の発電量は発電所1が100.0 MWであり、発電所2が50.0 MWであることが確認できました。

問題

２つの投資先があり、投資先１の期待リターンは10％、リスクが5％、投資先２の期待リターンは15％、リスクが8％とします。

また、ポートフォリオの全体的なリスクは6％未満にしたいと考えています。

この２つの投資先に関して、最適なポートフォリオを計算してください。

解き方・ソースコード

PythonのPuLPライブラリを使用してこの問題を解いてみます。

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

# 必要なライブラリをインポートする
import pulp

# 投資先とそれぞれの期待リターンとリスクを定義する
assets = ['Asset 1', 'Asset 2']
returns = {'Asset 1': 0.1, 'Asset 2': 0.15}
risks = {'Asset 1': 0.05, 'Asset 2': 0.08}

# 最適化問題を定義する
prob = pulp.LpProblem('Portfolio Optimization', pulp.LpMaximize)

# 投資先の割合を決定する変数を作成する
weights = pulp.LpVariable.dicts('Weight', assets, lowBound=0, upBound=1)

# 目的関数を設定する
prob += pulp.lpSum([weights[a] * returns[a] for a in assets])

# 制約条件を設定する
prob += pulp.lpSum([weights[a] for a in assets]) == 1
prob += pulp.lpSum([weights[a] * risks[a] for a in assets]) <= 0.06

# 最適化問題を解く
prob.solve()

# 結果を表示する
print('Optimal Portfolio:')
for a in assets:
    print('{}: {:.2%}'.format(a, weights[a].value()))

print('Expected Return: {:.2%}'.format(pulp.value(prob.objective)))
print('Risk: {:.2%}'.format(pulp.lpSum([weights[a] * risks[a] for a in assets]).value()))

このソースコードでは、投資先の割合を決定する変数を作成し、目的関数を最大化するように設定し、制約条件を設定しています。

最後に、prob.solve()を呼び出して問題を解決し、結果を表示しています。

[実行結果]

Optimal Portfolio:

Asset 1: 66.67%

Asset 2: 33.33%

Expected Return: 11.67%

Risk: 6.00%

投資先１に66.67％、投資先２に33.33％を割り当てると最適なポートフォリオとなり、期待リターンは11.67％、リスクは6.00％であることが確認できました。

問題

$1$ から $N$ までの整数が一個ずつ書かれた $ N \times N $ のマス目 $P$ が与えられます。

次の２種類の操作を繰り返すことで、すべての $k$ に対して「整数 $k$ が上から $k$ 行目・左から $k$ 列目のマスに存在する」ようにするためには、最小何回の操作が必要ですか。

🔹隣接する２つの行を交換する。
🔹隣接する２つの列を交換する。

[制約]
🔹$ 2 \leqq N \leqq 100 $

解き方・ソースコード

まず行の操作と列の操作を分解できないかどうかを考えます。

$i$ 行目に書かれた唯一の整数を $X_i$、$j$ 列目に書かれた唯一の整数を $Y_j$ とすると、以下の２点が成り立ちます。

🔹$i$行目と $i+1$ 行目を交換した場合： $X_i$ と $X_{i+1}$ のみが入れ替わり、$Y$ は不変
🔹$j$列目と $j+1$ 列目目を交換した場合： $Y_j$ と $Y_{j+1}$ のみが入れ替わり、$X$ は不変

このことは「行に対する操作」と「列に対する操作」を分けてよいことを意味しています。

また、目的通りの盤面になっていることと、$X=[1,2,3,4]$ かつ $Y=[1,2,3,4]$ になっていることはまったく同じです。

したがって、入力例１の最小操作回数は

🔹値１：最小何回の「隣接要素の交換」で、配列 $X$ を $[1,2,3,4]$ にできるか。
🔹値２：最小何回の「隣接要素の交換」で、配列 $Y$ を $[1,2,3,4]$ にできるか。

の総和となります。

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

#-------- 入力例１ ---------
N = 4                       # 整数の数
P = [
    [0, 0, 2, 0],
    [3, 0, 0, 0],
    [0, 0, 0, 4],
    [0, 1, 0, 0],
]
#---------------------------
# 「交換数を求める問題」2 つに分解する
X = [ None ] * N
Y = [ None ] * N
for i in range(N):
    for j in range(N):
        if P[i][j] != 0:
            X[i] = P[i][j]
            Y[j] = P[i][j]

# 交換数を求める関数
def inversion(A):
    answer = 0
    for i in range(len(A)):
        for j in range(i + 1, len(A)):
            if A[i] > A[j]:
                answer += 1
    return answer

# 答えを出力
print('解：', inversion(X) + inversion(Y))

[実行結果]

解： 5

最小操作回数が５回であることが確認できました。

PuLPとは

PuLPはPythonで書かれたオープンソースの線形および整数計画問題を解くためのモデリングライブラリです。

PuLPは、Pythonで線形最適化問題を解決するための高水準のインターフェースを提供し、制約条件、目的関数、変数を自然な形式で定義できます。

PuLPは、商用および非商用のプロジェクトで使用できます。

PuLPを使用することで、複雑な最適化問題を簡単に定式化および解決できるようになります。

これにより、様々な業界や分野で問題解決の効率が向上し、効果的な意思決定が可能になります。

問題

以下は、、PuLPを使用して解くことができる簡単な整数計画問題のサンプルです。

この問題では、変数 $x$ と $y$ を定義し、以下の制約条件を満たしながら目的関数を最大化することが目的です。

【目的関数】
🔹$ 3x + 4y $

【制約条件】
🔹$ 2x + y \leqq 20 $
🔹$ 4x - 5y \geqq -10 $
🔹$ x + 2y \leqq 14 $
🔹$x, y$は非負の整数である

ソースコード

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

import pulp

# 問題の定義
problem = pulp.LpProblem('Simple Integer Programming Problem', pulp.LpMaximize)

# 変数の定義
x = pulp.LpVariable('x', lowBound=0, cat='Integer')
y = pulp.LpVariable('y', lowBound=0, cat='Integer')

# 目的関数の定義
problem += 3*x + 4*y, 'Objective Function'

# 制約条件の定義
problem += 2*x + y <= 20, 'Constraint 1'
problem += 4*x - 5*y >= -10, 'Constraint 2'
problem += x + 2*y <= 14, 'Constraint 3'

# 最適化の実行
problem.solve()

# 結果の出力
print(f'Status: {pulp.LpStatus[problem.status]}')
print(f'Optimal Solution: {pulp.value(problem.objective)}')
print(f'x: {pulp.value(x)}')
print(f'y: {pulp.value(y)}')

まず、problemオブジェクトを作成します。

このオブジェクトは、数理最適化問題のモデルを表します。

ここでは、問題名と最大化を指定しています。

次に、変数$x$と変数$y$を定義します。

ここでは、下限が0の整数変数であることを指定しています。

目的関数は、$ problem += 3 * x + 4 * y $ と定義されています。

この問題では、この目的関数を最大化することが目的です。

制約条件は、それぞれ以下のように定義されています。

🔹制約1: $ 2 * x + y <= 20 $
🔹制約2: $ 4 * x - 5 * y >= -10 $
🔹制約3: $ x + 2 * y <= 14 $

最後に、problem.solve()を呼び出すことで、問題を解きます。

pulp.LpStatus[problem.status]を使用すると、解の状態を表示することができます。

また、pulp.value()を使用することで、最適解の目的関数の値、変数$x$と$y$の値を表示することができます。

[実行結果]

Status: Optimal
Optimal Solution: 36.0
x: 8.0
y: 3.0

最適解は、x=8, y=3であり、目的関数の値は36となっています。

制約条件を満たしながら、目的関数の最大値を得ることができました。

問題

配列 $A=[A_1, A_2, … , A_N]$ が与えられます。

大小関係を崩さないように、配列をできるだけ圧縮してください。

ここでの圧縮とは、以下の条件を全て満たす配列 $B=[B_1, B_2, … , B_N]$ を求める操作です。

なお、このような配列 $B$ は1通りに決まります。

🔹条件１ $B_1,B_2, … , B_N$ は1以上の整数である。
🔹条件２ $A_i < A_j$ であるような組 $(i,j)$ については、$B_i<B_j$ である。
🔹条件３ $A_i = A_j$ であるような組 $(i,j)$ については、$B_i=B_j$ である。
🔹条件４ $A_i > A_j$ であるような組 $(i,j)$ については、$B_i>B_j$ である。
🔹条件５ 条件１～条件４を満たす中で、配列 $B$ の最大値をできるだけ小さくする。

[制約]
🔹$ 1 \leqq N \leqq 100000 $
🔹$ 1 \leqq A_i \leqq 10^9 $

解き方

この問題はソートして重複を消す方法で解くことができます。

手順は以下の通りです。

１．配列 $A$ を昇順にソートします。
２．ソート後の配列 $A$ について、重複を除いた配列 $C$ を作成します。
３．配列 $C$ の要素に対して、$C_i$ が元の配列 $A$ の何番目の要素であったかを示す配列 $P$ を作成します。
　　つまり、$A_{P_i} = C_i$ となるような配列 $P$ を作成します。
４．配列 $B$ を用意し、$B_i$ を $P_i$ 番目の要素とすることで、配列 $B$ を求めます。

この方法で求められる配列 $B$ も、与えられた条件を満たす圧縮された配列となります。

ソースコード

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

#-------- 入力例１ ---------
N = 5                       # 整数の数
A = [46, 80, 11, 36, 46]    # 配列A
#---------------------------
import bisect

# 配列 T の作成（重複も消す）
T = list(set(A))
T.sort()

# 答えを求める
B = [None] * N
for i in range(N):
    B[i] = bisect.bisect_left(T, A[i])
    B[i] += 1

# 答え
print("解：", [str(i) for i in B])

このコードは、配列 $A$ をできるだけ圧縮するために、配列 $T$ を作成し、各要素に対応する圧縮された値を求めるものです。

まず、配列 $T$ を作成するために、set()関数を使い、重複を消した後、list()関数でリスト化しています。

その後、sort()メソッドを使い、要素を昇順に並べ替えています。

この $T$ 配列は、$A$ 配列を圧縮するための目印となるもので、条件２と３において、大小関係を崩さずに圧縮することができます。

次に、None で初期化した $B$ 配列を作成します。その後、for ループを用いて、配列 $A$ の各要素に対応する圧縮された値を求めます。

この圧縮された値は、配列 $T$ の中で、$A_i$ が何番目に小さいかを示しています。

bisect_left() 関数を使うことで、$A_i$ が $T$ 配列の中でどの位置にあるかを求めています。

bisect_left() 関数は、二分探索を行い、挿入するべきインデックスを返す関数です。

ここで、返される値は、$0$ から始まるため、$+1$ する必要があります。

最後に、$B$ 配列を表示して終了となります。

このコードは、時間計算量 $O(N \log N)$ で、与えられた条件を満たす圧縮された配列 $B$ を求めています。

[実行結果（入力例１）]

解： ['3', '4', '1', '2', '3']

問題

$N$ 個のタンクと $M$ 本のパイプがあります。

$j$ 本目のパイプはタンク $A_j$ からタンク $B_j$ の方向に毎秒 $C_j$ リットルまでの水を流すことができます。

タンク $1$ からタンク $N$ まで毎秒最大何リットルの水を流すことができますか。

ただし、タンクに水をためることはできません。

[制約]
🔹$ 2 \leqq N \leqq 400 $
🔹$ 1 \leqq M \leqq 4000 $
🔹$ 0 \leqq C_j \leqq 5000 $
🔹$ 答えは5000以下の整数 $

解き方

この問題は最大流問題と呼ばれる、グラフ上の最大流量を求める問題の一種です。

辺の容量はそのパイプが運べる最大の水の量 $C_j$ とします。

このグラフに対して、最大流量を求めるアルゴリズムを実行することで、タンク $1$ からタンク $N$ への最大流量を求めることができます。

有名な最大流アルゴリズムには、Ford-Fulkerson アルゴリズムや、Dinic アルゴリズムがあります。

ソースコード

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

#--------- 入力例１ ----------
N = 6
M = 7
A = [1, 1, 2, 2, 3, 4, 5]
B = [2, 4, 3, 5, 6, 5, 6]
C = [5, 4, 4, 7, 3, 3, 5]
#-----------------------------

from collections import defaultdict

class FordFulkerson:
    def __init__(self, n):
        self.n = n
        self.graph = defaultdict(list)
        self.visited = [False] * n
    
    def add_edge(self, u, v, c):
        self.graph[u].append([v, c, len(self.graph[v])])
        self.graph[v].append([u, 0, len(self.graph[u]) - 1])
    
    def dfs(self, v, t, f):
        if v == t:
            return f
        self.visited[v] = True
        for i, (nv, cap, rev) in enumerate(self.graph[v]):
            if self.visited[nv] or cap == 0:
                continue
            d = self.dfs(nv, t, min(f, cap))
            if d == 0:
                continue
            self.graph[v][i][1] -= d
            self.graph[nv][rev][1] += d
            return d
        return 0
    
    def max_flow(self, s, t):
        flow = 0
        while True:
            self.visited = [False] * self.n
            f = self.dfs(s, t, float('inf'))
            if f == 0:
                break
            flow += f
        return flow

max_flow = FordFulkerson(N + 1)
for i in range(M):
    max_flow.add_edge(A[i], B[i], C[i])

print(max_flow.max_flow(1, N))

このソースコードは、最大フロー問題を解くためにフォード・ファルカーソン法を実装したものです。

フォード・ファルカーソン法は、エッジのフロー量を増加させることによって最大フローを求めるアルゴリズムです。

このコードでは、最大フロー問題を解くために、FordFulkersonというクラスを定義しています。

このクラスは、グラフの構造を保持するための変数graphと、最大フローを求めるための関数max_flowを持ちます。

まず、入力例１のデータに従ってグラフを構築します。

このグラフは、defaultdictを用いて辞書形式で表現されており、各頂点にはその頂点に接続するエッジ（辺）の情報が保持されています。

次に、最大フローを求めるために、dfs関数を定義しています。

この関数は、sからtへの最大フローを求めるために、深さ優先探索を実行します。

具体的には、現在の頂点から次の頂点へのパスを探索し、最大フローを更新していきます。

最後に、max_flow関数を呼び出すことで、sからtへの最大フローを求めることができます。

この関数は、whileループを用いて、dfs関数を反復的に呼び出して最大フローを求めます。

[実行結果（入力例１）]

8

入力例１の場合、タンク $1$ からタンク $N$ へ毎秒最大 $8$ リットルの水を流すことができることが分かりました。

問題

頂点数 $n$ の無向グラフがあります。

隣接している頂点同士が違う色になるように、頂点に色を塗っていきます。

２色以内ですべての頂点を塗ることができるかどうか判定してください。

多重辺やループはないもとのします。

[制約]
🔹$ 1 \leqq N \leqq 1000 $

解き方

この問題は、二部グラフかどうかを判定する問題として知られています。

二部グラフとは、グラフを2つの部分集合に分け、同じ部分集合に属する頂点同士は隣接しないようにできるグラフのことです。

以下の手順に従って、二部グラフかどうかを判定することができます。

１．任意の頂点を1つ選び、そこからBFS（幅優先探索）またはDFS（深さ優先探索）を行います。
２．隣接する頂点のうち、まだ色を塗っていない頂点に、異なる色を塗ります。
３．2で塗られた頂点をスタートにして、再度BFSまたはDFSを行い、色を塗っていない頂点に色を塗ります。
４．2と3を繰り返し行い、すべての頂点に色を塗ることができたら、そのグラフは二部グラフであり、2色以内で塗ることができます。
　　一方、どのようにしてもすべての頂点に色を塗ることができなければ、そのグラフは二部グラフではありません。

ソースコード

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

#-------- 入力例１ ---------
n = 3            # 頂点数
# 隣接リスト
adj_list = [
    (1, 2),     # 頂点0につながっている頂点のリスト
    (0, 2),     # 頂点1につながっている頂点のリスト
    (0, 1)      # 頂点2につながっている頂点のリスト
]
#-------- 入力例２ ---------
n = 4            # 頂点数
# 隣接リスト
adj_list = [
    (1, 3),     # 頂点0につながっている頂点のリスト
    (0, 2),     # 頂点1につながっている頂点のリスト
    (1, 3),     # 頂点2につながっている頂点のリスト
    (0, 2)      # 頂点3につながっている頂点のリスト
]
#---------------------------
def is_bipartite(adj_list):
    color = [-1] * len(adj_list)
    color[0] = 0
    queue = [0]
    while queue:
        v = queue.pop(0)
        for u in adj_list[v]:
            if color[u] == -1:
                color[u] = 1 - color[v]
                queue.append(u)
            elif color[u] == color[v]:
                return False
    return True

if is_bipartite(adj_list):
    print("Yes")
else:
    print("No")

このコードでは、is_bipartite関数で、グラフが二部グラフかどうかを判定しています。

具体的には、1つ目の頂点を0番目の色で塗り、BFSで隣接する頂点に異なる色を塗っていくことで、二部グラフであるかどうかを判定しています。

二部グラフである場合は、すべての頂点に色を塗ることができるので、結果として“Yes”を出力します。

一方、二部グラフでない場合は、すべての頂点に色を塗ることができないため、結果として“No”を出力します。

[実行結果（入力例１）]

No

入力例１のグラフを塗るためには３色必要となるため、解は“No”となります。

[実行結果（入力例２）]

Yes

入力例２のグラフは、２色で塗ることができるため、解は“Yes”となります。

問題

頂点数 $N$、辺数 $M$ のグラフがあります。

頂点には $1$ から $N$ までの番号が付けられており、辺 $i$ は頂点 $A_i$ と $B_i$ を双方向に結ぶ長さ $C_i$ の辺です。

このグラフの最大全域木における辺の長さの総和を求めて下さい。

[制約]
🔹$ 2 \leqq N \leqq 100000 $
🔹$ 1 \leqq M \leqq 100000 $
🔹$ 1 \leqq C_i \leqq 10000 $
🔹このグラフは連結である。

解き方

全域木とは、$M$ 個の辺の中からいくつかを選んで作った全ての頂点が繋がっている木のことです。

全域木の中でも「長さの合計」が最大となるものを最大全域木といいます。

最大全域木は『長い辺から追加していく』という単純な貪欲法によって導き出すことができます。

ソースコード解説

前回記事では最小全域木の辺の長さの総和を求めましたが、今回の最大全域木の辺の長さの総和を求める場合は、辺の長さのソート順を逆にする変更だけで対応できます。

ソースコードの50行目だけ変更しており、辺の長い順にソートしています。

その他の処理は前回記事同様ですので、処理詳細はそちらをご参照下さい。

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

#-------- 入力例１ ---------
N = 7       # 頂点数
M = 9       # 辺数
# つながっている２つの頂点とその長さ
edges = [
    (1, 2, 12),
    (1, 3, 10),
    (2, 6, 160),
    (2, 7, 15),
    (3, 4, 1),
    (3, 5, 4),
    (4, 5, 3),
    (4, 6, 120),
    (6, 7, 14)
]
#---------------------------
# Union-Find 木
class unionfind:
    # n 頂点の Union-Find 木を作成
    def __init__(self, n):
        self.n = n
        self.par = [ -1 ] * (n + 1) # 最初は親が無い
        self.size = [ 1 ] * (n + 1) # 最初はグループの頂点数が 1
    
    # 頂点 x の根を返す関数
    def root(self, x):
        # 1 個先（親）がなくなるまで進み続ける
        while self.par[x] != -1:
            x = self.par[x]
        return x
    
    # 要素 u, v を統合する関数
    def unite(self, u, v):
        rootu = self.root(u)
        rootv = self.root(v)
        if rootu != rootv:
            # u と v が異なるグループのときのみ処理を行う
            if self.size[rootu] < self.size[rootv]:
                self.par[rootu] = rootv
                self.size[rootv] += self.size[rootu]
            else:
                self.par[rootv] = rootu
                self.size[rootu] += self.size[rootv]
    
    #  要素 u と v が同一のグループかどうかを返す関数
    def same(self, u, v):
        return self.root(u) == self.root(v)

# 辺を長さの大きい順にソート
edges.sort(key = lambda x: -x[2])

# 最小全域木を求める
uf = unionfind(N)
answer = 0
for a, b, c in edges:
    if not uf.same(a, b):
        uf.unite(a, b)
        answer += c

# 答えの出力
print('解：', answer)

[実行結果（入力例１）]

解： 321

最大全域木における辺の長さの総和を求めることができました。

問題

頂点数 $N$、辺数 $M$ のグラフがあります。

頂点には $1$ から $N$ までの番号が付けられており、辺 $i$ は頂点 $A_i$ と $B_i$ を双方向に結ぶ長さ $C_i$ の辺です。

このグラフの最小全域木における辺の長さの総和を求めて下さい。

[制約]
🔹$ 2 \leqq N \leqq 100000 $
🔹$ 1 \leqq M \leqq 100000 $
🔹$ 1 \leqq C_i \leqq 10000 $
🔹このグラフは連結である。

解き方

全域木とは、$M$ 個の辺の中からいくつかを選んで作った全ての頂点が繋がっている木のことです。

全域木の中でも「長さの合計」が最小となるものを最小全域木といいます。

最小全域木は『短い辺から追加していく』という単純な貪欲法によって導き出すことができます。

この方法はクラスカル法と呼ばれており、配列のソートとUnion-Find木を使って実装することができます。

ソースコード解説

下記のソースコードは、Union-Find 木を使って与えられたグラフの連結成分を管理し、クエリーを処理するプログラムです。

入力例１では、3つの頂点と4つのクエリーが与えられています。

クエリー１は頂点同士を結ぶ辺を追加し、クエリー２は２つの頂点が同じ連結成分に属するかを判定するものです。

まず、Union-Find 木のクラスを定義しています。

コンストラクタでは、頂点数 n を受け取り、par と size の初期化を行っています。

par は親の頂点を保持する配列で、初期値は全て -1 です。

size はグループの頂点数を保持する配列で、初期値は全て 1 です。

次に、root 関数では、引数で指定された頂点 x の根を探索して返します。

par[x] が -1 になるまで、頂点 x の親に移動していきます。

unite 関数は、引数で指定された要素 u, v を同一のグループに統合します。

この関数内では、それぞれの根を取得し、もし根が異なる場合はグループの頂点数が小さい方を大きい方に統合します。

par 配列には親の頂点を、size 配列にはグループの頂点数を記録しています。

same 関数は、引数で指定された要素 u, v が同一のグループに属するかどうかを判定します。

この関数内では、root 関数を呼び出して u, v の根を比較します。

次にクエリーの処理を行います。

クエリー１の場合は、union 関数を呼び出して頂点 u, v を統合します。

クエリー２の場合は、same 関数を呼び出して頂点 u, v が同一のグループに属するかどうかを判定します。

最後に最小全域木における辺の長さの総和を出力しています。

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

#-------- 入力例１ ---------
N = 7       # 頂点数
M = 9       # 辺数
# つながっている２つの頂点とその長さ
edges = [
    (1, 2, 12),
    (1, 3, 10),
    (2, 6, 160),
    (2, 7, 15),
    (3, 4, 1),
    (3, 5, 4),
    (4, 5, 3),
    (4, 6, 120),
    (6, 7, 14)
]
#---------------------------
# Union-Find 木
class unionfind:
    # n 頂点の Union-Find 木を作成
    def __init__(self, n):
        self.n = n
        self.par = [ -1 ] * (n + 1) # 最初は親が無い
        self.size = [ 1 ] * (n + 1) # 最初はグループの頂点数が 1
    
    # 頂点 x の根を返す関数
    def root(self, x):
        # 1 個先（親）がなくなるまで進み続ける
        while self.par[x] != -1:
            x = self.par[x]
        return x
    
    # 要素 u, v を統合する関数
    def unite(self, u, v):
        rootu = self.root(u)
        rootv = self.root(v)
        if rootu != rootv:
            # u と v が異なるグループのときのみ処理を行う
            if self.size[rootu] < self.size[rootv]:
                self.par[rootu] = rootv
                self.size[rootv] += self.size[rootu]
            else:
                self.par[rootv] = rootu
                self.size[rootu] += self.size[rootv]
    
    #  要素 u と v が同一のグループかどうかを返す関数
    def same(self, u, v):
        return self.root(u) == self.root(v)

# 辺を長さの小さい順にソート
edges.sort(key = lambda x: x[2])

# 最小全域木を求める
uf = unionfind(N)
answer = 0
for a, b, c in edges:
    if not uf.same(a, b):
        uf.unite(a, b)
        answer += c

# 答えの出力
print('解：', answer)

[実行結果（入力例１）]

解： 55

最小全域木における辺の長さの総和を求めることができました。

問題

ある地域には $N$ 個の駅と $M$ 本の鉄道路線があります。

駅には $1$ から $N$ までの番号が付けられていて、$i$ 本目の路線は駅 $A_i$ と駅 $B_i$ を双方向に結んでいます。

この地域に台風が上陸すると、いくつかの路線は運休になってしまいます。

以下のようなクエリ―が $Q$ 個与えられるのでそのクエリ―を処理してください。

　クエリー１：$x$ 本目の路線が運休になる。
　クエリー２：運休状況を踏まえて駅 $s$ から駅 $t$ へ移動できるかどうかを答える。

解き方・ソースコード

この問題も前回記事同様、Union-Find木を使って解いていきますが、辺を減らしながらの連結状況を答えるのは難しいので、クエリーを逆向きにして辺を増やしながら処理することにします。

まずソースコードの冒頭では、次のように入力データを設定しています。

　・N： 頂点の数（駅の数）。
　・M： 辺の数（２つの駅を結ぶ路線の数）。
　・edges： 辺のリスト。各要素は、頂点の番号のペアで表される。
　・Q： クエリ―の数。
　・query： クエリ―のリスト。各要素は、1または2の数字と、いくつかの頂点の番号の組み合わせで構成される。

このプログラムは、以下のような処理を行います。

　１．辺のリストをインデックス0から始まるように修正する。
　２．クエリ―に応じて、最後まで残っている辺のリストを取得する。
　３．Union-Find木を用いて、最後まで残っている辺でグラフを構成する。
　４．クエリ―を逆順に処理し、連結を増やしながら移動可能かどうかの回答を求める。
　５．回答を出力する。

Union-Find木を使ってグラフを構成していきます。

Union-Find木は、グループを管理するためのデータ構造です。

各要素は、自身が属するグループの根の番号を覚えています。

グループをマージするときには、根の番号を統一します。

これによって、任意の2つの要素が同じグループに属するかどうかを簡単に判定することができます。

このプログラムは、最後まで残っている辺でグラフを構成します。

そして、クエリ―を逆順に処理しながら、Union-Find木を使ってグループをマージしていきます。

最後に、各クエリ―に対する回答を出力します。

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

#-------- 入力例１ ---------
N = 5       # 駅の数
M = 5       # ２つの駅をつなぐ路線の数
edges = [
    (1, 2),
    (2, 3),
    (3, 4),
    (4, 5),
    (3, 5)
]

Q = 5       # クエリ―数
query = [
    [2, 1, 5],
    [1, 3],
    [2, 1, 5],
    [1, 5],
    [2, 1, 5]
]
#---------------------------
edge = []
# インデックスを0スタートにする
for A, B in edges:
    edge.append((A - 1, B - 1))

# 最後まで残っている辺を求める
last = [True] * M
for q in query:
    if q[0] == 1:
        q[1] -= 1
        last[q[1]] = False
    else:
        q[1] -= 1
        q[2] -= 1

# union-find木
uf = [-1] * N
def root(i):
    while True:
        if uf[i] < 0:
            return i
        if uf[uf[i]] < 0:
            return uf[i]
        uf[i] = uf[uf[i]]
        i = uf[i]

def unite(i, j):
    i = root(i)
    j = root(j)
    if i == j:
        return
    # union by size
    if uf[i] > uf[j]:
        i, j = j, i
    uf[i] += uf[j]
    uf[j] = i

ans = []
for i in range(M):
    if last[i]:
        A, B = edge[i]
        unite(A, B)

# クエリ―を逆順に処理する
for q in reversed(query):
    if q[0] == 1:
        A, B = edge[q[1]]
        unite(A, B)
    else:
        _, U, V = q
        ans.append("移動可能" if root(U) == root(V) else "移動不可")

# クエリ―の順番を戻して、回答を表示
for s in reversed(ans):
    print(s)

[実行結果（入力例１）]

移動可能

移動可能

移動不可

クエリ―に応じて、駅間の移動ができるかどうかを出力することができました。