ABC072 C - Together Python での解答

Question

ABC072 C - Together
この問題について自分のコードだとテストサンプルのうち１つだけWAになってしまいます．どこが間違っているのかを教えてほしいです．

コード

import numpy as np

n = int(input())
a = np.array(list(map(int, input().split())))

a = np.sort(a)

max_num = 0
p = 0

while n - p >= max_num:
    num = 0
    a -= a[p]
    while p < n and a[p] <= 2:
        num += 1
        p += 1
    if num > max_num:
        max_num = num

print(max_num)

各変数の説明
n : 数列の数
a : 数列
max_num : a_i = x となる数の最大値でこの問題の返り値．
p : 現在着目しているaのインデックス．詳細は後述．
num : 各ループごとのa_i = x となる数．

アルゴリズム

ある数を基準にしてその数を数列のすべての数から引いたとき，0, 1, 2 となる要素の数が操作後にa_i = x となる数だからその最大値を探すというのが基本的な考え方です．

1. n, a のインプット．後で要素全体の減算を行うためndarrayを使用．事前に降順にソートしておく．
2. 以下whileループ内の説明．まず数列のすべての数から数列の最小値の値を引きます．はじめはa[p]=a[0]が最小値なのでa[p]の値を引いています．
3. ２つ目のwhileループ内の説明．2の操作後，aの先頭から要素の値が0, 1, 2である数をカウントし，numに保存します．ここでpは着目しているaのインデックスを表していてa[p]が0, 1, 2 であればインデックスを一つ進め，numを１増やします．a[p]>2となるまで同様に繰り返し，a[p]>2となったタイミングでのnumがそのループでのa_i = x となる数です．
4. numがmax_num よりも大きかったらmax_num を更新し，2, 3を繰り返します．最終的なmax_num が答えとなります．

Answer 1

7
0 1 2 3 3 3 3

のようなテストケースを考えてみてください。
すると while p < n and a[p] <= 2:している間にp = 3となってしまいます。
実際はp = 1の所でnumがどれくらいになるかをカウントしないといけないのですが、それができていないのでWAとなります。
正解は6 (p = 1)ですが、投稿者様のコードだと4 (p = 3)になっちゃいます。

この問題は

n = int(input())
a = list(map(int, input().split())
count = [0 * 100001]

for num in a:
    count[a] += 1

として、各数字が何回現れるかをカウントアップします。
それからcount[i] + count[i + 1] + count[i + 2]が最大となる`iをみつければそれが答えです。