🧰

# キーエンスプログラミングコンテスト2021 B - Mex Boxes

2021/01/19に公開

AtCoder

キーエンスプログラミングコンテスト2021 B - Mex Boxes

問題へのリンク

問題概要

$N$ 個の一つの整数が書かれたボールがある。ボールに書かれている数字は $a_1, a_2, \dots, a_N$ である。
この $N$ 個のボールを $K$ 個の箱に割り振って入れる。どのボールも箱に入れるが、ボールが入ってない箱があってもよい。
すべてのボールを入れた後にそれぞれの箱の蓋に数字が $x$ が表示される。 $x$ は箱の中に $x$ が書かれたボールが存在しないような最小の非負整数である。このとき表示される数の総和の最大値を求めよ。

制約

$1 \leq K \leq N \leq 3 × 10^5$
$1 \leq N \leq 2 × 10^5$
$0 \leq a_i \leq N$

ABC中の解答

同じ箱に0, 1, 2, ...と連続するようにボールを入れていけば蓋に表示される数の総和が大きくなる。したがって、同じ箱に0や1などの小さい値のボールを入れずに1個ずつ散らすように入れればよい。そこで $a_1, a_2, \dots, a_N$ に含まれる値をカウントしていく。途中で数字が切れた場合はそれ以上の値は箱に入れてないと考えても影響しない。ここらへんまで考えて実装を始めた。カウントにははじめ numpy.unique を用いようと思ったが途中数字が切れたかどうかの判定が大変だと思ったので collections.Counter を使うのに途中方針転換した。

実装を始めるときには 数字が切れた場合はそれ以上の値は考えない と頭にあったのに実装中にぬけて1WA
K 個と箱の制約で得られる値似上限(例えば0のボールが10個あっても箱が3個しかなければ3個しか考えなくてよい)があるのを忘れて1WA
の計2WAをして少しもったいなかった。

累積Minをうまく考えられたのはよかったかな。

from collections import Counter
N, K = map(int, input().split())
A = list(map(int, input().split()))

counts = Counter(A)
for i in range(1, len(counts)):
    counts[i] = min(counts[i - 1], counts[i])

MAX = 3 * 10**5 + 1
ans = 0
v = min(K, counts[0])
for i in range(1, MAX):
    c = counts.get(i, None)
    if c is None:
        ans += v * i
        break
    if v > c:
        ans += (v - c) * i
        v = c
print(ans)

maspyさんの解法

公式の解説ではコードがなかったのでmaspyさんの解法を眺めてみたら numpy.bincount が使われていてなるほどと思った。 numpy.unique ではなく numpy.bincount ならボールがなかった数字も調べやすい。その後も numpy.minimum.accumulate を使い累積Minを求めていた。 numpy の使い方が上手だなと改めて思った。見習わないと。
総和に影響をあたえるボールつまり0, 1, ... と連続しているボールしかないとみなすと ボールの数 = 箱の蓋の数字の総和 になるので最後に numpy.sum で総和を求めているのもうまいと思った。

maspyさんのをぱっとみた後に自分でも勉強のため実装したコードを下に示す。

import numpy as np

N, K = map(int, input().split())
A = list(map(int, input().split()))

counts = np.bincount(A)
counts = np.minimum(counts, K)
counts = np.minimum.accumulate(counts)
ans = counts.sum()
print(ans)

cf. maspyさんの解答