身長だけの場合
仮に、身長についてだけ考えればよいとすれば、割と簡単な問題です。事前にソートした上で、いわゆる「尺取り法」で実装できます。
A.sort() head = tail = 0 a_max = a_min = A[head] max_len = 0 while True: while a_max - a_min > K: tail += 1 a_min = A[tail] while a_max - a_min <= K: max_len = max(max_len, head - tail + 1) head += 1 if head == N: print(max_len) return a_max = A[head]
体重を含めた場合
上記の尺取り法の中では、「身長の観点で許容されるグループ」が自然に網羅されています。許容されるグループの中で、人数が最大になるものを検索する形になります。
そこで、「身長の観点で許容されるグループ」のそれぞれに対して、そのメンバーから「体重の観点でも許容されるサブグループ」を構成するという方法が考えられます。
極端な方法としては、「身長の観点で許容されるグループ」のそれぞれに対して、体重についての尺取り法を実行するというやり方が考えられます。
上記の実装では、tail を固定した状態で、head を伸ばしていき、最大限に伸ばし切ったタイミングで、体重についての尺取り法を呼び出すという実装になっています。いくつか TLE していますが、方向性としては悪くなさそうです。
TLE の 1 つの要因として、体重についての尺取り法を呼び出す際に、対象となるサブグループのコピーが発生する点がちょっと重い気がします。
そこで、体重に関しては、配列 num_B[i] を用意して、今考えている「身長の観点で許容されるグループ」の中で、
・num_B[i] = 体重が i 〜 i + K の範囲の人数
をトラッキングする様にします。こうすれば、max(num_B) として、身長・体重の両方について許容される最大人数を得ることができます。max(B) を計算するタイミングを最小限になるようにうまく調整すると無事に AC します。