CodeIQ:行列をうまく整理して!

私自身が表題の問題を解いた時のプログラムについて解説します。
問題の詳細は「行列をうまく整理して!」(CodeIQ)を参照してください。

問題の概要

問題を引用します。
ここでは、そんな行列をうまく整理するために、列を区切る方法を考えてみます。
n人が並んでいる行列を一人ずつに区切ることにします。
ただ、一つの列を一度に区切ることができるのは一人だけです。
区切る列が3つある場合は、同時に3人で区切ることができます。

店員さんが最大m人いるとき、最短何回で区切ることができるかを求めてください。
例えば、n=8, m=3のときは次の図のようになり、4回で区切ることができます。

   □□□□□□□□
1回目 □□□□|□□□□
2回目 □□|□□|□□|□□
3回目 □|□|□|□
4回目       □|□

標準入力から、nとmの値がコンマ区切りで指定されます。最短で区切るための回数を標準出力に出力してください。

私のプログラム

Pythonで解答しています。

#!/usr/local/bin/python3

import fileinput
import math

Results = []
Cnt = 0

def part(n,m):
	global Cnt
	for r in Results:
		nxt = []
		p = 0

		for i in r:
			if (p<m):
				if i//2 > 1 : nxt.append(i//2)
				if i - i//2 > 1 :nxt.append(i-i//2)
				p += 1
			else:
				nxt.append(i)

		nxt = sorted(nxt, reverse = True)
		Cnt += 1
		Results.append(nxt)

		if not nxt:
			break

if __name__ == "__main__":
	for line in fileinput.input():
		line = line.strip()
		if not line:
			continue

		n,m = map(int, line.split(","))
		Results.append([n])
		part(n,m)

		print(Cnt)

解説

それではコードの説明をします。

考え方

問題のとおり、素朴に分割します。
つまり、1回あたり最大で店員の人数に等しい列数までを分割することができ、列数が店員の人数を超えた場合、超えた分は次の回に分割するという処理を繰り返します。
管理する必要があるのは現在何列になっているのかと、各列に現在何人いるのかという情報です。実際のコードではResultsでこの情報を管理しています。

Results

Resultsは分割したお客の列情報の履歴を管理します。実際には履歴にする必要はなく、最新の情報だけあれば良いのですが、デバッグのため履歴を保持しています。
Resultsは次のような構造をとります。例のようにお客の数が8人の場合を考えます。

[初期状態]
Results[0] = [8]
初期状態(つまり0回目の分割)では8人の列が1つだけという状態を持ちます。

[1回目]
Results[0] = [8]
Results[1] = [4,4]
1回分割されたので要素番号1に[4,4]という値を持ちます。
このリストの要素数が列の数、各値が列の人数になります。

[2回目]
Results[0] = [8]
Results[1] = [4,4]
Results[2] = [2,2,2,2]
同様に分割されてこのようになります。

[3回目]
Results[0] = [8]
Results[1] = [4,4]
Results[2] = [2,2,2,2]
Results[3] = [2]
Results[2]の先頭から3要素が分割されてそれぞれ[1,1]になりますが、人数が1になった行列はそれ以上分割できないので除いてしまいます。そのため、Result[3]のようになります。

[3回目]
Results[0] = [8]
Results[1] = [4,4]
Results[2] = [2,2,2,2]
Results[3] = [2]
Results[4] = []
最後の列が分割され、1人の列は除かれるのでこのようになります。

part()

前述のResultsの操作を行っているのがこの関数です。
11行目のループでResultsから要素を取り出して処理します。Resultsには最初1つしか要素がなく、分割の結果新たな要素が追加された場合は25行目で要素数が増えるのでこのループは無限ループのように動作します。
変数nxtは分割結果を保持し、pは分割のために動いた店員の数を管理します。

15行目のループは最大で店員の数までの行列を分割する処理を行います。16行目のifの条件が店員の数をチェックし、店員の数を超えた場合20行目のelseで行列を処理せずそのままnxtに追加して次の処理に回します。
17〜18行目の条件は1人になった列を除くためにあります。

重要なのが23行目のソートで、行列を人数に従って降順にソートします。これによって人数が多い行列から処理されることになるのですが、この処理がないと余分な手順がかかってしまいます。

24行目で何回目の処理かをカウントし、25行目で次の分割のためResultsの要素を追加します。全て1人ずつに分割されたらnxtが空になるので27行目でそれをチェックしてループを終了します。

雑感

やってみたらなかなか面白かったのですが、ちょっと問題がわかりにくくかなり長い間放置してあった問題の1つです。
例を見ながら少し考えればわかる話ではあるのですが、1回の分割で余った店員の扱いが問題のルールになかったのでその部分でちょっと仕様が不明確でした。つまり、例の場合で行くと1回目は店員が1人で分割しますが、2回目は残りの3人しか使えないのか、リセットされてもう一度4人なのかということです。