Pythonで配列のペアの合計がkで割り切れるかどうかを確認するプログラム

問題概要

偶数個の要素を含む配列 nums と整数 k が与えられたとします。この配列をちょうど n/2 組のペアに分割し、それぞれのペアの合計が k で割り切れるようにできるかを判定するのが課題です。条件を満たす組み合わせが存在すれば True を、そうでなければ False を返します。

たとえば、nums = [9,5,3,4,7,10,20,8]、k = 3 の場合を考えてみましょう。(9, 3)、(5, 7)、(4, 20)、(8, 10) というペアを作ることができ、すべてのペアの合計が3で割り切れるため、出力は True になります。

解決手順

この問題は、各要素を「kで割った余り」の観点で扱うことで解けます。具体的には次の手順に従います。

• 空のリスト dp と、初期値0のカウンター count を用意します。
• 配列 nums の各要素 x について、以下を繰り返します。
  • t = k - (x % k) を計算します。
  • t == k(つまり x が k で割り切れる)の場合は、count を1増やします。
  • それ以外の場合は、t をリスト dp の末尾に追加します。
• count % 2 != 0(kで割り切れる要素が奇数個)の場合は False を返します。これらの要素は互いにペアを組む必要があるためです。
• リスト dp を昇順にソートします。
• ポインタ low = 0、high = len(dp) - 1 を設定します。
• low < high の間、以下を繰り返します。
  • dp[low] + dp[high] != k であれば False を返します。
  • low を1増やし、high を1減らします。
• すべてのペアが条件を満たしていれば True を返します。

実装例

以下の実装を見ると、理解がより深まるでしょう。

def solve(nums, k):
   dp=[]
   count=0
   for x in nums:
      t=k-(x % k)
      if t == k:
         count+=1
      else:
         dp.append(t)
   if count % 2 != 0:
      return False
   dp.sort()
   low = 0
   high = len(dp)-1
   while low < high:
      if dp[low] + dp[high] != k:
         return False
      low += 1
      high -= 1
   return True

nums = [9,5,3,4,7,10,20,8]
k = 3
print(solve(nums, k))

入力

[9,5,3,4,7,10,20,8], 3

出力

True

アルゴリズムのポイント

要素 x と y の合計が k で割り切れるのは、(x mod k) + (y mod k) が k または 0 になるときと同じです。そこで、余りが0の要素は互いにペアを組めるため個数だけを数え、それ以外の要素については「k − 余り」という相方の値をリスト dp に記録します。ソート後、両端から中央へ向かって2つのポインタを動かしながら合計が k になるかを確認すれば、全要素が適切な相手と対応しているかを O(n log n) の計算量で効率的に検証できます。