Pythonで3要素の合計がターゲット未満となるトリプレットの個数を数えるプログラム

解くための手順は以下の通りです。

1. リスト nums を昇順にソートする
2. ans := 0 と初期化する
3. n := nums のサイズとする
4. i を 0 から n−1 までループする
   • k := n − 1 とする
   • j を i+1 から n−1 までループする
     • k > j かつ nums[i] + nums[k] + nums[j] >= target の間、k を 1 減らし続ける
     • j == k になったらループを抜ける
     • ans += k − j を加算する
5. ans を返す

このアルゴリズムのポイントは、リストがソート済みであるため、j と k の間に含まれるすべての要素 m について nums[i] + nums[j] + nums[m] < target が自動的に成り立つという点です。つまり、条件を満たす範囲の幅（k − j）を一度に数え上げることができ、個別の組み合わせをすべて列挙する必要がありません。

実装例

以下にPythonでの実装例を示します。

class Solution:
   def solve(self, nums, target):
      nums.sort()
      ans = 0
      n = len(nums)
      for i in range(n):
         k = n - 1
         for j in range(i + 1, n):
            while k > j and nums[i] + nums[k] + nums[j] >= target:
               k -= 1
            if j == k:
               break
            ans += k - j
      return ans

ob1 = Solution()
nums = [-2, 6, 4, 3, 8]
target = 12
print(ob1.solve(nums, target))

入力

[-2, 6, 4, 3, 8], 12

出力

5

まとめ

このように、ソートと2ポインタ技法を活用することで、3重ループによる総当たりよりも大幅に効率的に、合計がターゲット未満となるトリプレットの個数を求めることができます。同様の手法は「3Sum」系の問題全般に応用できるため、ぜひ覚えておきましょう。