Pythonで最小値と最大値の和がk以下となる部分列の個数を求めるプログラム

問題の概要

配列 nums と整数 k が与えられたとき、nums の空でない部分列のうち、「部分列内の最小要素と最大要素の和が k 以下」という条件を満たすものの個数を求めます。

答えは非常に大きな値になる可能性があるため、10⁹ + 7 で割った余りを返します。

入力例

nums = [4, 6, 7, 8]、k = 11 の場合、出力は 4 になります。条件を満たす部分列は次の通りです。

• [4] … 最小値 4、最大値 4 → 4 + 4 ≤ 11
• [4, 6] … 最小値 4、最大値 6 → 4 + 6 ≤ 11
• [4, 6, 7] … 最小値 4、最大値 7 → 4 + 7 ≤ 11
• [4, 7] … 最小値 4、最大値 7 → 4 + 7 ≤ 11

解法のアプローチ:ソート + ツーポインタ

この問題は、配列をソートした上で両端から動かす2つのポインタ(ツーポインタ法)を使うことで効率的に解けます。手順は以下の通りです。

1. リスト nums を昇順にソートする
2. m := 10⁹ + 7(剰余演算用の定数)
3. left := 0(左端ポインタ)
4. right := 配列サイズ − 1(右端ポインタ)
5. res := 0(結果格納用変数)
6. left ≤ right の間、以下を繰り返す
   • nums[left] + nums[right] > k の場合:right を 1 減らす
   • それ以外の場合:
     • num_inside := right − left(間にある要素数)
     • res := (res + 2^num_inside) mod m
     • left を 1 増やす
7. res を返す

なぜ 2^num_inside なのか?

ソート済みの配列では nums[left] が最小値になります。nums[left] + nums[right] ≤ k が成り立つとき、left の要素を必ず含み、left+1 〜 right の範囲の各要素については「含める / 含めない」を自由に選べます。その組み合わせは 2^{(right−left)} 通りあり、これを left ごとに足し合わせることで重複なくすべての部分列を数えられます。

実装例(Python)

def solve(nums, k):
    nums.sort()
    m = 10**9 + 7
    left = 0
    right = len(nums) - 1
    res = 0
    while(left <= right):
        if nums[left] + nums[right] > k:
            right -= 1
        else:
            num_inside = right - left
            res = (res + pow(2, num_inside, m)) % m
            left += 1
    return res

nums = [4,6,7,8]
k = 11
print(solve(nums, k))

入力

[4,6,7,8], 11

出力

4

計算量について

ソートに O(n log n)、その後のツーポインタ走査は各要素が一度しか処理されないため O(n) となります。全体の計算量は O(n log n) であり、全ての部分列を列挙する O(2ⁿ) の素朴な方法と比べて大幅に高速です。