Pythonで線形和割り当て問題を解く

昔、あるアルゴリズムを実装する中で使ったことがある、 linear_sum_assignment っていうscipyのメソッドを久々に使うと思ったら使い方を忘れていたのでその復習を兼ねた記事です。

これは、2部グラフの最小重みマッチングとも呼ばれている問題で、要するに、二つのグループの要素からそれぞれ1個ずつ選んだペアにコストが定義されていて、どのように組み合わせてペアを選んでいったらコストの和を最小にできるかという問題です。

この説明はわかりにくいですね。もう少し具体的なのがいいと思うので、Scipyのドキュメントで使われている例を使いましょう。

Scipyのドキュメントではworker(作業者)とjob(仕事)を例に解説されています。
参考: scipy.optimize.linear_sum_assignment — SciPy v1.9.1 マニュアル

例えば、4人の作業者がいて4つの仕事があったとします。そして、その4人がそれぞれの仕事をした場合に、かかる時間(=コスト)が次のように与えられていたとします。行列形式ですが、i行j列の値が、作業者iが仕事jを実行した場合にかかるコストです。(例を乱数で作りました。)

import numpy as np


np.random.seed(0)
cost = np.random.randint(1, 10, size=(4, 4))
print(cost)

"""
[[6 1 4 4]
 [8 4 6 3]
 [5 8 7 9]
 [9 2 7 8]]
"""

cost[1, 2] = 6 ですが、これは作業者1が仕事2を行った場合のコストが6ということです。
(インデックスが0始まりであることに注意してください。cost[1, 2]は2行3列目の要素です。)

さて、上の図を見ての通り、作業者ごとに仕事の得手不得手があり、コストが違うようです。そこで、これらの仕事をそれぞれ誰が担当したらコストの総和を最小にできるでしょうか、というのが線形和割り当て問題です。

これが、先ほどの linear_sum_assignment を使うと一発で解けます。

ドキュメントにある通り、戻り値が行のインデックス、列のインデックスと帰ってくるので注意してください。

from scipy.optimize import linear_sum_assignment


row_ind, col_ind = linear_sum_assignment(cost)
print("行:", row_ind)
print("列:", col_ind)
"""
行: [0 1 2 3]
列: [2 3 0 1]
"""

二つのarray(プリントしてるのでlistに見えますがnumpyのArrayです)が戻ってきます。
これが、worker0がjob2を担当し、worker1がjob3を担当し、、、と読んでいきます。
これがコストを最小にする組み合わせです。簡単でしたね。

さて、値の戻ってき方がちょっと独特だったのでプログラムでこれを使うにはコツが要ります。こう使うと便利だよ、ってところまでドキュメントに書いてあると嬉しいのですが、書いてないので自分で考えないといけません。

インデックスとして返ってきているので、次のようにコスト行列のインデックスにこの値を入れると、最適化された組み合わせのコストが得られます。そして、sum()すると合計が得られます。以下の通り、14が最小ということがわかります。

print(cost[row_ind, col_ind].sum())
# 14

Scipyの実装を疑うわけではないのですが、念の為、本当にこの組み合わせが最適で14が最小なのか、全組み合わせ見ておきましょう。itertools.permutationsを使います。

from itertools import permutations


for perm in permutations(range(4)):
    print(list(perm), "=>", cost[range(4), perm].sum())
"""
[0, 1, 2, 3] => 25
[0, 1, 3, 2] => 26
[0, 2, 1, 3] => 28
[0, 2, 3, 1] => 23
[0, 3, 1, 2] => 24
[0, 3, 2, 1] => 18
[1, 0, 2, 3] => 24
[1, 0, 3, 2] => 25
[1, 2, 0, 3] => 20
[1, 2, 3, 0] => 25
[1, 3, 0, 2] => 16
[1, 3, 2, 0] => 20
[2, 0, 1, 3] => 28
[2, 0, 3, 1] => 23
[2, 1, 0, 3] => 21
[2, 1, 3, 0] => 26
[2, 3, 0, 1] => 14
[2, 3, 1, 0] => 24
[3, 0, 1, 2] => 27
[3, 0, 2, 1] => 21
[3, 1, 0, 2] => 20
[3, 1, 2, 0] => 24
[3, 2, 0, 1] => 17
[3, 2, 1, 0] => 27
"""

どうやらあってそうですね。

col_ind の方を使って、行列を並び替えることもできます。i行目のworkerがi列目のjobを担当する直感的に見やすい行列が次のようにして得られます。

print(cost[:, col_ind])
"""
[[4 4 6 1]
 [6 3 8 4]
 [7 9 5 8]
 [7 8 9 2]]
"""

また、解きたい問題や実装によっては、この行と列の対応を辞書にしたほうが使いやすいこともあるでしょう。そのような時はdictとzipで変換します。

print(dict(zip(row_ind, col_ind)))
# {0: 2, 1: 3, 2: 0, 3: 1}

ここまでの例では、与えられた行列はコストの行列でこれを最小化したい、という問題設定でやってきました。ただ、場合によっては利益やスコアの行列が与えられて、最大化する組み合わせを探したいという場合もあると思います。行列にマイナス掛けて同じことすればいいのですが、linear_sum_assignment自体にもそれに対応した引数があります。

それが、maximize で、 デフォルトはFalseですが、Trueにすると最大化を目指すようになります。同じ行列でやってみます。さっき全パターン列挙しているので正解はわかっていて、[0, 2, 1, 3]か[2, 0, 1, 3]のどちらかが得られるはずです。

print(linear_sum_assignment(cost, maximize=True))
# (array([0, 1, 2, 3]), array([0, 2, 1, 3]))

[0, 2, 1, 3]の方が出ててきましたね。

ここまで、正方行列を取り上げてきましたが、linear_sum_assignment は、一般行列についても実行できます。行と列の数が違う場合は、行と列のうち数が小さい方に揃えて、実行されます。

まず、行が多い(workerが多い)場合をやってみましょう。7行4列で、7人のworkerがいて、jobが4つあって、コストがそれぞれ定義されていた場合に、どの4人を選抜してそれぞれにどの4つのタスクをやってもらうのが最適か、という問題を解くのと対応します。

np.random.seed(0)
cost = np.random.randint(1, 10, size=(7, 4))
print(cost)
"""
[[6 1 4 4]
 [8 4 6 3]
 [5 8 7 9]
 [9 2 7 8]
 [8 9 2 6]
 [9 5 4 1]
 [4 6 1 3]]
"""

row_ind, col_ind = linear_sum_assignment(cost)
print("行:", row_ind)
print("列:", col_ind)
"""
行: [0 2 5 6]
列: [1 0 3 2]
"""

次に同様に横長の行列の場合です。例えば4人のworkerがいて7つのjobがあったときに、どの4つのjobを選んで実行したら利益を最大化できるか、って問題がこれに相当します。(最小化でいい例が思いつかなかったのでこれは最大化でやります。)

np.random.seed(0)
score = np.random.randint(1, 10, size=(4, 7))
print(score)
"""
[[6 1 4 4 8 4 6]
 [3 5 8 7 9 9 2]
 [7 8 8 9 2 6 9]
 [5 4 1 4 6 1 3]]
"""

row_ind, col_ind = linear_sum_assignment(score, maximize=True)
print("行:", row_ind)
print("列:", col_ind)
"""
行: [0 1 2 3]
列: [4 5 3 0]
"""

以上が linear_sum_assignment の基本的な使い方になります。

コメントを残す

メールアドレスが公開されることはありません。