Pythonでソート済み配列が合計kのペアに分割できるか判定する方法

数値の配列ともう一つの数値 k が与えられたとき、その配列を「各ペアの合計が必ず k になるように」ペアへ分割できるかどうかを判定する問題を考えます。

たとえば、入力が arr = [1, 2, 3, 4, 5, 6]、k = 7 の場合、出力は True になります。(2, 5)、(1, 6)、(3, 4) というペアを作れば、それぞれの合計がすべて 7 になるからです。

アルゴリズムの考え方

配列がすでにソートされていることを利用すると、両端から中央に向かって進める「二ポインタ（Two Pointers）」手法で効率的に判定できます。手順は以下の通りです。

• n を配列の要素数とします。
• n が奇数の場合、ペアを組み切れないため False を返します。
• low = 0、high = n - 1 として初期化します。
• low < high の間、以下を繰り返します。
  • arr[low] + arr[high] が k と等しくなければ False を返します。
  • low を 1 増やし、high を 1 減らします。
• ループを抜けたら True を返します。

実装例

def solve(arr, k):
    n = len(arr)
    if n % 2 == 1:
        return False
    low = 0
    high = n - 1
    while low < high:
        if arr[low] + arr[high] != k:
            return False
        low = low + 1
        high = high - 1
    return True

arr = [1, 2, 3, 4, 5, 6]
k = 7
print(solve(arr, k))

入力

[1, 2, 3, 4, 5, 6], 7

出力

True

計算量について

このアルゴリズムは配列を一度だけ走査するため、時間計算量は O(n) です。また、追加のメモリをほとんど使わないため、空間計算量は O(1) で済みます。

なお、この手法は配列がソート済みであることが前提です。未ソートの配列に対しては、あらかじめソートを行っておく必要があり、その場合は全体の計算量が O(n log n) になります。