【SRM757Div1】CentipedeSocks【250point】

問題概要

$F$本の足があるムカデが$C$匹いる. そのムカデたちに靴下を履かせたい. 但し, それぞれのムカデに履かせる靴下の色は全て同じでないといけない(ムカデ同士はどうでもいい).

今, ビンの中にいくつかの靴下がある. その情報は$sockCount$として渡され, 色$C_1$の靴下は$sockCount_1$個, 色$C_2$の靴下は$sockCount_2$個, …である.

ビンから無作為に$X$個の靴下を取ったとき, それが必ずムカデたちの要望を満たすような$X$の最小値を求めなさい.

制約

• $1 \leq C \leq 50$
• $1 \leq F \leq 100$
• $1 \leq |sockCounts| \leq 100$
• $1 \leq sockCounts_i \leq 10^7$

考察

最悪の場合を考えれば良さそう.

$sockCount$の各要素について, $F$未満だったら, 1匹の要望も満たせないので全部取る.

以下では, $F$以上の場合を考える.

1. とりあえず1匹の要望も満たさないように, $F - 1$足ずつ取る.
2. 引いた後の値について, さらに$p$($1 \leq p \leq F$)個取れば1匹の要望を満たせる.
   • 最悪の場合を考えるので, なるべく$p$が大きくなるように取るのが最適.
   • 優先度付きキューで最大値から取り出せるようにすれば良い.
3. $C$回ループを回して, popして$p$個取ってpushしてを繰り返す.
   • 途中で足りなくなったら$-1$を返す(下の実装だと最初に走査したときにダメだったら$-1$して逃げてるけど).

だけど…

ここまで実装したけどサンプルが合わない.

サンプル2で39を返していたところに注目すると, 最後の1匹は特別視しなければならないことに気づく. (最後の1匹は2足以上取ると冗長なので)

結局$C - 1$回ループを回して, 最後にres += 1した.

実装

Pythonの優先度付きキューであるheapqの使い方を学んだ.

最小値からしか取り出せないの不便.

最大値から取り出したかったので$-1$倍してpushして, popして$-1$倍した(13, 17, 22行目).

(これからはなんかデータ構造使いたいときは素直にC++使おう…)

提出コード(Python3🐍)

from heapq import *

class CentipedeSocks:
    def fewestSocks(self, C, F, sockCounts):
        ret, cnt = 0, 0
        pq = []
        for sock in sockCounts:
            cnt += sock // F
            if sock < F:
                ret += sock
            else:
                ret += F - 1
                heappush(pq, -(sock - (F - 1)))
        if cnt < C:
            return -1
        for i in range(C - 1):
            sock = -heappop(pq)
            if sock < F:
                ret += sock
            else:
                ret += F
                heappush(pq, -(sock - F))
        # 最後の1匹
        ret += 1
        return ret