Toyota Programming Contest 2024#1（AtCoder Beginner Contest 337）F問題メモ

Toyota Programming Contest 2024#1（AtCoder Beginner Contest 337）

GよりむずいF問題

F - Usual Color Ball Problems

• $N$ 個のボールが1列に並び、それぞれの色は $C_i$
• $M$ 個の、1つにつきボールを $K$ 個まで入れられる空の箱がある
• $i=1,2,...,N$ について、以下を求めよ
  • ボールを、$i$ 番目が先頭になるように並べなおす。$1,...,i-1$ 番目のボールはそのままの順番で末尾にもっていく
  • 先頭から順に、以下の要領でボールを箱に入れる（または捨てる）
    • ボールの色を $c$ として、既に色 $c$ が入っていて $K$ 個未満の箱があれば、そこに入れる
    • なければ、空の箱があれば適当に1つ選んで入れる
    • それもなければ捨てる
  • 最終的に、箱に入ったボールの個数を求めよ
• $1 \le M,K \le N \le 2 \times 10^5$

$i=1$ の時を求めて、1つずらすだけだとほとんどは同じ状態になるだろうので、$i=2,3,...$ の答えは差分で計算したい。

しかし、実際に箱に入るボールは列中で飛び飛びに現れるので、 区間を1つずらしたときに、どの色の箱が無くなり、どのボールが新たな箱を作るのか、結構変化してしまい、効率的に管理しづらい。。

M=3 K=2

Ci         1 2 3 5 2 5 4
箱に入る?  o o o x o x x

Ci           2 3 5 2 5 4 1
箱に入る?    o o o o o x x

Ci             3 5 2 5 4 1 2
箱に入る?      o o o o x x o

Ci               5 2 5 4 1 2 3
箱に入る?        o o o o x o x

重要な考察として、最終的に色 $c$ のボールが箱に入る個数というのは、以下の小さい方になる。
（これに気付かないと進めないのに、なかなか気付きづらい……）

• 色 $c$ のボールが入った箱数 $\times K$
• 色 $c$ のボールの個数 $g(c)$

なので、先頭からボールを入れていくようにシミュレーションする際、 「全ての箱の色が確定したタイミング」、つまり「最後の箱に、1個ボールが入ったタイミング」で、 全体の箱に入ったボールの個数が確定するので、打ち切ってよい。

そこまでのボールは、必ず何かしらの箱に入っていることが保証される。
飛び飛びでは無くなるので、$i$ をずらすたび、途中から探索を行える。

Ci         1 2 3 5 2 5 4
箱に入る?  o o o - - - -      ^ : 最後の箱にボールが入ったタイミング
               ^              - : 未調査
Ci           2 3 5 2 5 4 1
箱に入る?    o o o - - - -
                 ^
Ci             3 5 2 5 4 1 2
箱に入る?      o o o - - - -
                   ^
Ci               5 2 5 4 1 2 3
箱に入る?        o o o o - - -
                       ^

尺取法の要領で、区間の左右端のポインタを管理しておき、区間を1つずらした時、

• 左端のボールを取り除く。取り除いた前後で箱が1つ減れば、差分更新する
• それによって箱が1つ空になった場合、右端からボールを追加していく。次に新たに箱が作られたら差分更新し、終わる

高々2色について、答えを更新すればよい。

Python3

from collections import Counter, defaultdict

n, m, k = map(int, input().split())
ccc = list(map(int, input().split()))

if k == 1:
    for _ in range(n):
        print(m)
    exit()

cnt = Counter(ccc)
boxes = defaultdict(int)
total_box = 0
i = 0
for i in range(n):
    c = ccc[i]
    boxes[c] += 1
    if boxes[c] % k == 1:
        total_box += 1
        if total_box == m:
            break
else:
    for _ in range(n):
        print(n)
    exit()


def ceil_k(x):
    return ((x - 1) // k + 1) * k


ans = []
tmp = 0
for c, t in boxes.items():
    tmp += min(ceil_k(t), cnt[c])
ans.append(tmp)

r = i + 1
ccc = ccc * 2

for l in range(n - 1):
    cl = ccc[l]
    boxes[cl] -= 1
    if boxes[cl] % k == 0:
        total_box -= 1
        tmp -= min(ceil_k(boxes[cl] + 1), cnt[cl])
        tmp += min(ceil_k(boxes[cl]), cnt[cl])

    while total_box < m:
        cr = ccc[r]
        boxes[cr] += 1
        if boxes[cr] % k == 1:
            total_box += 1
            tmp -= min(ceil_k(boxes[cr] - 1), cnt[cr])
            tmp += min(ceil_k(boxes[cr]), cnt[cr])
        r += 1

    ans.append(tmp)

print('\n'.join(map(str, ans)))