Pythonで異なる単語がちょうどk個含まれる部分リストの数を求めるプログラム

単語のリストと値 k が与えられたとき、異なる単語がちょうど k 個含まれる部分リスト（連続する要素からなる部分配列）の個数を求める問題を考えます。

例えば、入力が words = ["Kolkata", "Delhi", "Delhi", "Kolkata"]、k = 2 の場合、出力は 5 になります。これは、次の5つの部分リストに2種類のユニークな単語が含まれているためです。

• ["Kolkata", "Delhi"]
• ["Delhi", "Kolkata"]
• ["Kolkata", "Delhi", "Delhi"]
• ["Delhi", "Delhi", "Kolkata"]
• ["Kolkata", "Delhi", "Delhi", "Kolkata"]

一方、["Delhi", "Delhi"] はユニークな単語が1種類しか含まれていないため、カウント対象外となります。

解法のアプローチ：スライディングウィンドウ

この問題は「スライディングウィンドウ（尺取り法）」を使うことで O(n) の時間計算量で効率的に解けます。

ポイントは、まず「異なる単語が k 個以下である部分リストの総数」を求める関数 work() を定義し、最後に work(words, k) - work(words, k - 1) を計算することで「ちょうど k 個」の場合の数を導き出す点です。

アルゴリズムの手順

• 引数として words と k を受け取る関数 work() を定義する
• n := words のサイズ
• k が 0 の場合は 0 を返す
• cnt := 各単語の出現回数を記録する新しい辞書（マップ）
• ans := 0、l := 0（ウィンドウの左端）
• r を 0 から n-1 までループ:
  • word := words[r]
  • word が cnt に存在しなければ cnt[word] := 0
  • cnt[word] を +1 する
  • cnt のサイズ（異なる単語数）が k を超えている間:
    • cnt[words[l]] を -1 する
    • cnt[words[l]] が 0 になったら、そのキーを cnt から削除する
    • l を +1 して左端を進める
  • ans に (r - l + 1) を加算（右端 r を含む有効なウィンドウの数）
• ans を返す

メイン処理では (work(words, k) - work(words, k - 1)) を返します。

Pythonでの実装例

理解を深めるために、実際のコードを見てみましょう。

class Solution:
   def solve(self, words, k):
      return self.work(words, k) - self.work(words, k - 1)

   def work(self, words, k):
      n = len(words)
      if k == 0:
         return 0
      cnt = dict()
      ans = 0
      l = 0
      for r in range(n):
         word = words[r]
         if word not in cnt:
            cnt[word] = 0
         cnt[word] += 1
         while len(cnt) > k:
            cnt[words[l]] -= 1
            if cnt[words[l]] == 0:
               del cnt[words[l]]
            l += 1
         ans += r - l + 1
      return ans

ob = Solution()
words = ["Kolkata", "Delhi", "Delhi", "Kolkata"]
k = 2
print(ob.solve(words, k))

入力

["Kolkata", "Delhi", "Delhi", "Kolkata"], 2

出力

5

計算量について

右端 r と左端 l はそれぞれ最大でも n 回しか移動しないため、時間計算量は O(n)、単語ごとの出現回数を保持する辞書が必要となるため空間計算量も O(n) となります。全ての部分リストを列挙する素朴な O(n²) 以上のアプローチと比べ、非常に効率的です。