【Python】数値リストから長さ3以上の等差数列を数えるプログラム

数値のリスト nums が与えられたとき、その中に含まれる「長さ3以上の連続する等差数列」の個数を求める問題を考えます。等差数列とは、隣り合う数同士の差（公差）がすべて等しい数列のことです。

例えば、入力が nums = [6, 8, 10, 12, 13, 14] の場合、出力は 4 になります。次の4つの等差数列が見つかるためです。

• [6, 8, 10]
• [8, 10, 12]
• [6, 8, 10, 12]
• [12, 13, 14]

解法のアプローチ

この問題は、リストを一度走査するだけで解くことができます。基本的なアイデアは、「同じ差が何回連続して現れたか」をカウントし、そのカウントから等差数列の総数を算出するというものです。

具体的には、以下の手順に従います。

1. count := 0、ans := 0 で初期化します。

2. i を 2 から nums のサイズ未満まで繰り返します。

   • nums[i] - nums[i-1] が nums[i-1] - nums[i-2] と等しい場合は、count を 1 増やします。
   • それ以外の場合は、ans に (count * (count + 1)) // 2 を加算し、count を 0 に戻します。

3. ループ終了後に count が 0 以外であれば、ans に (count * (count + 1)) // 2 を加算します。

4. ans を返します。

なぜ count × (count + 1) ÷ 2 なのか？

同じ公差が count 回連続して一致した区間では、そこから切り出せる長さ3以上の等差数列は、長さ3のものが count 個、長さ4のものが count - 1 個…と続き、合計は 1 + 2 + … + count = count × (count + 1) ÷ 2 となります。この性質を使うことで、二重ループなしに効率的に答えを求められます。

実装例

それでは、実際のコードを見てみましょう。

class Solution:
    def solve(self, nums):
        count = 0
        ans = 0
        for i in range(2, len(nums)):
            if nums[i] - nums[i - 1] == nums[i - 1] - nums[i - 2]:
                count += 1
            else:
                ans += (count * (count + 1)) // 2
                count = 0
        if count:
            ans += (count * (count + 1)) // 2
        return ans

ob = Solution()
nums = [6, 8, 10, 12, 13, 14]
print(ob.solve(nums))

入力

[6, 8, 10, 12, 13, 14]

出力

4

このアルゴリズムの時間計算量は O(n)、空間計算量は O(1) であり、リストの要素数が大きくなっても高速に動作するのが特徴です。