AtCoder青色のkym3535と申します。昨日のARC133のC問題がなかなか面白かったので、解説記事を書いてみたいと思います。アルゴリズムや典型知識ではなく、純粋に考察力が問われる問題で、自分にとってはかなり好みなタイプでした。
問題
(AtCoder Regular Contest 133) C - Row Column Sums
H 行 W 列からなるマス目があります。
すぬけくんは、各マスに 0 以上 K−1 以下の整数を書き込もうとしています。ここで、以下の条件を満たす必要があります。
- 各 1<=i<=H について、 i 行目にあるマスに書かれた整数の総和を K で割った余りは A_i である。
- 各 1<=i<=W について、 i 列目にあるマスに書かれた整数の総和を K で割った余りは B_i である。
条件を満たすようにマスに整数を書き込むことができるかどうか判定し、また可能な場合は、書き込む整数の総和としてありうる最大値を求めてください。
考察
状況把握と明らかな「不可能」解
まずは入力例1で与えられている状況を簡単に図にしてみると、下の図のようになる。
1行目に書き込んだ4つの整数の和のmod3が0、2行目に書き込んだ4つの整数の和のmod3が2、という具合だ。
ここで一つ明らかなことがあって、Aの総和(この図では $0+2=2$ )と、Bの総和(この図では $1+2+2+0 = 5$)、それぞれのmod3が異なるなら、そのような条件を満たす書き込み方は存在しない。
なぜなら、Aの総和のmod3は、すべてのマスの総和のmod3と等しい。(ここで0は1行目の4つの数字の総和のmod3、2 は2行目の4つの数字の総和のmod3である。よって、0+2のmod3をとると、それは、すべてのマスの数字の総和のmod3と等しい。)
Bの総和も同様に、すべてのマスの数字の総和のmod3と等しくなる。
よって、この2つが異なるなら、矛盾が生じるため、条件を満たす書き込み方は無いとわかる。
単純に書き込んでみる
この例では、Aの総和は2、Bの総和は5と異なっている。ここで、行や列の和のmod3ではなく、行や列の和そのものを考えることにする。A・Bは、行や列で数字を合計した後に、mod3 をとっていることを考えると、1列目の和は0ではなく、K( 例1ではK=3 )を足した、3でもよいことがわかる。
試してみると、たとえば以下のような解が見つかる。
だが、これは書き込んだ数字の和を最大化できていない。
書き込む数字の和の最大化
先ほど1行目の和の「0」を「3」に書き換えたように、各行の和、列の和はK(この例では3)を足してもかまわない。
しかし、際限なく足してもいいわけでは無く、限度がある。1行の数字の和の最大値を考えると、マスに書き込める数字は0,1,2のみである。なので、ある行の4つのマスすべてに「2」が書かれているとき、その行の総和が最大となり、そのときの値は$2\times4=8$となる。
同様に、1列の総和の最大値は、$2\times2=4$である。
これに注意して、各行・列の和を同時に3ずつ増やしてみる。すると、次のような経過をたどる。
2行目の和の「5」は、3を足して「8」にすることが出来るが、あいにく、どの列の和も、3を足すと限度の「4」を超えてしまう。よって、これ以上、行・列の和を増やすことが出来ない。
数字を書き込めることを確かめる
実際に、行・列の和がこのようになる書き込み方を探してみる。
まず、列の和が最も大きい、1列目に注目する。各マスには2以下の数しか書き込めないので、1列目の2つのマスにはそれぞれ2を書き込む。
次に、列の和が2番目に大きい4列目に注目する。ここでは、各行にあとどれだけ足すとよいか(図では「残○」で表記)が出来るだけ平たくなるように足すことにする。そうすると、1行目に2、2行目に1を書き込むことになる。
あとは、2列目・3列目も同様に書き込むと、最終的には以下のようになる。
このような方法を用いると、各マスに書き込むべき数字を決めることが出来そうだとわかる。
厳密には、この方法が本当に正しいのか証明する必要がある。しかし、ぱっと考えてもなかなか反例が浮かばず、時間の無駄になると判断し、本番ではここからすぐに実装の考察に進んだ。
実装の考察
考察で見たようにAやBの各要素に何度もKを足すのは時間がかかりそうなので、もう少し効率的な方法を考えることにした。
すなわち、行と列で1ステップずつ、同時にKを足していくのではなく、行と列でそれぞれ、最大でどこまで大きく出来るかを調べることにした。この例では、まず
- 行： A = [0 2] に対し、8を超えない範囲で3を何度も足すと、[6 8] になる。
- 列： B = [1 2 2 0]に対し、4を超えない範囲で3を何度も足すと[4 2 2 3]になる。
を求める。計算は、次のようにすれば、何度も繰り返す必要が無い。A=[0 2]から[6 8]を求める部分を例にすると、
- 限度の 8 から A[0]=0 を引く。 8-0 = 8 なので、あと 8 まで足すとが出来る。
- 8//3 = 2 なので、3をあと2回足せる。
- A[0] に 3*2 = 6 を足して、A[0] は 0+6 = 6 になる。
- 同様に 8-A[1] = 6、6//3 = 2 なので、A[1]に 3*2 = 6 を足して、A[1] は 2+6 = 8 となる。
この結果のAの総和は14、Bの総和は11なので、このうち小さい方(Bの11)を答えとして採用する。
実装
解き方さえわかれば、実装については特に難しいことは無いと思います。本番の提出コードは下記の通りです。
本番の解答 【提出 #28693164】
オーダーは $O(H+W)$で、和をとったり足し算したりするだけなので、Pythonでも余裕で間に合うと判断しました。
本番ではA問題に5分、B問題に約55分、この問題に約15分かかり、1300点75分でした。
B問題に時間がかかりすぎたのが響き、いまいちパフォーマンスが伸びきらず残念…。
import sys
def input():
return sys.stdin.readline()[:-1]
h,w,k = tuple(map(int, input().split()))
a = list(map(int, input().split()))
b = list(map(int, input().split()))
if sum(a)%k != sum(b)%k :
print(-1)
exit()
a_lim = (k-1) * w
b_lim = (k-1) * h
a_sum = sum(a)
b_sum = sum(b)
for i in range(h):
add_n = (a_lim - a[i])//k
a_sum += add_n * k
for j in range(w):
add_n = (b_lim - b[j])//k
b_sum += add_n * k
print(min(a_sum,b_sum))