Pythonで1つの削除で出現頻度が揃う最長シーケンスを求めるプログラム

問題の概要

数値のリストが与えられたとき、「シーケンスから1つの数値を削除すると、残りのすべての数値が同じ回数だけ出現する」という条件を満たす、最長のシーケンスの長さを求める問題を考えてみましょう。

たとえば、入力が numbers = [2, 4, 4, 7, 7, 6, 6] の場合、出力は 7 になります。これは、先頭の 2 を削除すれば [4, 4, 7, 7, 6, 6] となり、4・7・6 がそれぞれ2回ずつ出現して条件を満たすためです。

解法のアプローチ

この問題は、リストを先頭から順に走査しながら、各時点での出現頻度の状態を効率よく管理することで解けます。まず、次のデータ構造を用意します。

• num_freq: 各数値の出現回数を記録する新しいマップ（辞書）
• freq_freq: 「出現回数ごとの数値の種類数」を記録する新しいマップ
• diff_freq: 現在存在する異なる出現回数を管理する新しいセット
• result: 結果を保持する変数（初期値は 1）

続いて、リスト nums の各インデックス i と各数値 num について、以下の処理を行います。

• cur_freq := num_freq[num] として、現在の num の出現回数を取得します。
• num_freq[num] を 1 増やします。
• freq_freq[cur_freq] を 1 減らし、freq_freq[cur_freq + 1] を 1 増やします。
• cur_freq + 1 を diff_freq に追加します。
• cur_freq が diff_freq に含まれており、かつ freq_freq[cur_freq] が 0 になっている場合は、cur_freq を diff_freq から削除します。
• diff_freq の要素から df_list という新しいリストを作成します。
• df_list のサイズが 1 の場合（すべての数値の出現回数が揃っている状態）、result := i + 1 とします。
• それ以外で、df_list のサイズが 2 であり、かつ次の2つの条件を both 満たす場合も result := i + 1 とします。
  • [|freq_freq[df_list[0]] − freq_freq[df_list[1]]|、freq_freq[df_list[0]]、freq_freq[df_list[1]]] のいずれかが 1 である
  • [|df_list[0] − df_list[1]|、df_list[0]、df_list[1]] のいずれかが 1 である

最後に result を返します。

条件分岐の意味

df_list のサイズが 2 のときの条件は、「あと1つの要素を削除するだけで全頻度を揃えられる状態」を判定しています。具体的には、次のようなケースが該当します。

• 片方の頻度を持つ数値がちょうど1つだけ存在する（その要素を削除すればよい）
• 頻度の差が 1 で、大きい方の頻度を持つ数値が1つだけ存在する（その要素を1つ削除すればよい）
• 頻度 1 の数値が存在する（その要素を丸ごと削除すればよい）

この方法では、各要素を1回ずつ処理するだけでよいため、全体の計算量は O(n) となり、大きな入力に対しても高速に動作します。

実装例

それでは、理解を深めるために以下の実装を見てみましょう。

from collections import defaultdict
class Solution:
   def solve(self, nums):
      num_freq = defaultdict(int)
      freq_freq = defaultdict(int)
      diff_freq = set()
      result = 1
      for i, num in enumerate(nums):
         cur_freq = num_freq[num]
         num_freq[num] += 1
         freq_freq[cur_freq] -= 1
         freq_freq[cur_freq + 1] += 1
         diff_freq.add(cur_freq + 1)
         if cur_freq in diff_freq and freq_freq[cur_freq] == 0:
            diff_freq.remove(cur_freq)
         df_list = list(diff_freq)
         if len(df_list) == 1:
            result = i + 1
         elif (
            len(df_list) == 2
            and any(
               x == 1
               for x in [
                  abs(freq_freq[df_list[0]] - freq_freq[df_list[1]]),
                  freq_freq[df_list[0]],
                  freq_freq[df_list[1]],
               ]
            )
            and any(x == 1 for x in [abs(df_list[0] - df_list[1]), df_list[0], df_list[1]])
         ):
            result = i + 1
      return result
ob = Solution()
print(ob.solve([2, 4, 4, 7, 7, 6, 6]))

入力

numbers = [2, 4, 4, 7, 7, 6, 6]

出力

7