Pythonでシフトを自動で組む

2015年にシフトを自動で組むプログラムをPythonで書いた - matsu7874のブログを書いているのですが、これはランダムで100案作って制約を満たすか確認する雑な実装なので、もう少し勉強して書き直してみるシリーズの第一回目です。

問題設定

N日分のシフト枠をM人のワーカーで埋める

入力

N,Mと各ワーカーのNG日程のリスト

例：5日3人の場合、下記の入力でworker0はday2がNG、worker3はday0,1,2がNGということを表します。

[
  [2,],
  [3,],
  [4,],
  [0,1,2,],
  [2,3,4,],
]

満たすべき制約

各ワーカーのNGを優先する。
毎日capacity人以下の割当を行う。

なるべく実現したいこと

毎日の割当はなるべくcapacity人にする。
ワーカーの割当回数はなるべく均等にする。

方針

まずはシンプルな貪欲法で実装したいと思います。

リストの先頭から働ける人に働いてもらう。
割当回数を均等化するために、割当日数の少ない人から割り当てていく。

最悪時は毎日全員が働けるかをチェックするため、NGが重ならない限りは割当が可能です。

実装

import sys
import random
import heapq


class Worker:
    def __init__(self, name, ng):
        self.name = name
        self.ng = ng

    def __str__(self):
        return str((self.name, str(self.ng)))


class Shift:
    def __init__(self, days, capacity):
        self.days = days
        self.capacity = capacity
        self.shift = [['UNASSIGNED']*capacity for i in range(days)]

    def assign(self, workers):
        hq = [(0, i) for i in range(len(workers))]
        used_hq = []
        for d in range(self.days):
            day_assigned_cnt = 0
            while hq and day_assigned_cnt < self.capacity:
                worker_assigned_cnt, worker_id = heapq.heappop(hq)
                if d in workers[worker_id].ng:
                    heapq.heappush(used_hq, (worker_assigned_cnt, worker_id))
                else:
                    self.shift[d][day_assigned_cnt] = workers[worker_id].name
                    day_assigned_cnt += 1
                    heapq.heappush(used_hq, (worker_assigned_cnt+1, worker_id))
            while used_hq:
                heapq.heappush(hq, heapq.heappop(used_hq))
            self.shift[d].sort()
        worker_assigned_cnt = [0 for i in range(len(workers))]
        for wa_cnt, worker_id in hq:
            worker_assigned_cnt[worker_id] = wa_cnt
        print(worker_assigned_cnt, file=sys.stderr)

    def __str__(self):
        return ',\n'.join([str((i, d)) for i, d in enumerate(self.shift)])


def main():
    n_days = 10
    n_workers = 5
    # 各workerはランダムで最大n_days/2日分までNGを出す。
    workers = [Worker('worker {:02}'.format(i), sorted(random.sample(
        range(n_days), random.randint(0, n_days//2)))) for i in range(n_workers)]
    print(*workers, sep='\n', file=sys.stderr)
    shift = Shift(n_days, 2)
    shift.assign(workers)
    for d, ws in enumerate(shift.shift):
        print(d, '\t'.join(ws), sep='\t')


if __name__ == "__main__":
    main()

実行結果

上記のプログラムを実行すると、乱数を使っているので同じ結果にはならないと思いますが、下記のような結果を得ることができます。

('worker 00', '[]')
('worker 01', '[5, 9]')
('worker 02', '[0, 9]')
('worker 03', '[1, 2, 4, 7, 9]')
('worker 04', '[1, 3, 7, 8, 9]')
[5, 4, 4, 4, 2]
0       worker 00       worker 01
1       worker 00       worker 02
2       worker 01       worker 04
3       worker 02       worker 03
4       worker 00       worker 04
5       worker 02       worker 03
6       worker 01       worker 03
7       worker 00       worker 01
8       worker 02       worker 03
9       UNASSIGNED      worker 00

4/5がNGを出しているday9以外は割当が行えていることがわかります。一方で各ワーカーの割当回数には2から5のばらつきがあり、NGを5箇所出しているworker03,04間で割当回数が2,4と差がある点は課題です。また、day0にworker00の代わりにworker04を割り当てることでより均等な割当が実現可能です。

最後に

ナース・スケジューリング：問題把握とモデリングを読んでいるのですが、いろいろなアルゴリズムが応用されていて非常に面白いです。次回以降、内容を紹介しながらプログラムを改善していきたいと思います。