解説を読んでこの天才解法は思いつかないのでは..と思っていたんですが、結局自分の考え方と全く一緒だったのでメモ。
こう考える
まず、この問題を以下のように言い換えましょう。(と、簡単に書いていますが、コンテスト中はこの言い換えはできませんでした)
- 1x1のn色のブロックをAi個ずつ持っています。 (各部署の人がAi人ずついます)
- なるべく大きなP行k列の長方形を作りたいです (k人のP個のプロジェクトを作りたいです)
- ただし、ある行に(離れていても)同じ色があってはいけません。(同じプロジェクトに同じ部署の人がいてはいけません)
- 問:作れる最大のPを求めてください
これを図にします。
$k=4$ときに、$P=6$が作れるかを判定したいとします。この際に、部署の人数をソートする必要はありません。上と下では、部署の人数が少しだけ違います。
さて、条件は行 = 横に同じ色があってはいけませんでした。こう考えると、各色のブロックは、P個以上絶対に使えないのが見えます。逆に言うと、P以下であれば適当に積んでもいいです。(図では各色を縦に並べていますが、横に同じ色がいないなら適当な場所で良いです。
これは解説の通り、
for(long long x : v) sum += min(x, md);
if(sum >= k * md)
を判定しているだけです。(各ブロックをPまでしか使えないとして、K*Pが埋められるか判定しています)
実装
本番の自分の考え方は結果として上記と一緒で明らかに遠回りなコードを書いてしまいました。
def do():
n, k = map(int, input().split())
dat = list(map(int, input().split()))
total = sum(dat)
def func(mid): # midグループ作れるかの判定
compk = 0 # 全グループに埋めたk(これを満たせばOK)
curkneed = mid # そのkを作るにはあとこの人数必要です
for i in range(n): #部署を順番に見ていく
if compk == k: return True
curdept = dat[i] # 今の部署の人数をcurdeptとする
if curkneed <= curdept: # もし、今作りたいkに人が足りているなら全員突っ込む
compk += 1 # この場合、そのkは完成
curdept -= curkneed # そのkに必要な人数使ったので部署に残っている人数
used = curkneed
curkneed = mid # 本来、次のkにはmid人必要なんだけど
# もし、「今の部署に当てたグループ」以外に充てられるならその分は充てる。
curkneed -= min(curdept, curkneed - used) # ただし、"curkneed-used"使ってしまうと、絶対に同じグループにその部署の人が入ってしまうのでminをとる
if compk == k: return True
continue
# 今作りたいグループにこの部署を全員充てる
curkneed -= curdept
return False
ng = 10 ** 20 + 10
while (abs(ok - ng) > 1):
mid = (ok + ng) // 2;
if (func(mid)):
ok = mid;
else:
ng = mid;
print(ok)
do()