Pythonで有効な等差数列となるクエリの数を効率的に求めるプログラム

この記事では、Pythonを使って「部分列が等差数列になっているか」を判定する複数のクエリを効率的に処理する方法を解説します。

問題の概要

数値のリスト nums と、クエリのリスト queries が与えられます。各クエリは [i, j] の形式で表され、「nums[i] から nums[j] まで（両端を含む）の部分列は等差数列か？」という問い合わせを意味します。最終的な目標は、true（等差数列である）を返すクエリの個数を求めることです。

入力例と出力例

たとえば、次のような入力が与えられたとします。

• nums = [2, 4, 6, 8, 7, 6, 5, 2]
• queries = [[3, 4], [0, 3], [2, 4]]

この場合の出力は 2 になります。理由は以下の通りです。

• [2, 4, 6, 8] は公差 2 の等差数列なので、クエリ [0, 3] は true
• [8, 7] は2要素の等差数列なので、クエリ [3, 4] も true
• [6, 8, 7] は隣接要素の差が一定ではないため、クエリ [2, 4] は false

解法のアプローチ

各クエリに対して毎回部分列を走査すると計算量が大きくなるため、前処理によって高速に答えられるようにします。手順は以下の通りです。

1. nums が空の場合は 0 を返します。
2. n を nums のサイズとします。
3. diff: 隣接要素同士の差 nums[i + 1] - nums[i] を格納したリストを作成します（長さ n - 1）。
4. rle: 長さ n - 1 のリストを 0 で初期化します。ここには「その位置で終わる、連続して同じ差が続く長さ」を記録します。
5. i を 0 から n - 2 まで順に処理します。
   • i > 0 かつ diff[i] == diff[i - 1] の場合：rle[i] = rle[i - 1] + 1
   • それ以外の場合：rle[i] = 1
6. ans を 0 で初期化し、各クエリ (i, j) について判定します。
   • i == j の場合：要素が1つだけなので常に等差数列として ans += 1
   • それ以外の場合：rle[j - 1] >= (j - i) が成り立てば 1 を加算
7. ans を返します。

なぜこの方法でうまくいくのか？

rle[j - 1] は「インデックス j - 1 で終わる、同じ差が連続している区間の長さ」を表します。区間 [i, j] が等差数列であるためには、diff 配列上の位置 i から j - 1 までの差がすべて同一である必要があります。つまり、連続する同じ差が少なくとも (j - i) 個続いていればよいので、rle[j - 1] >= (j - i) という条件で判定できるのです。これにより、各クエリを O(1) で処理できます。

Pythonでの実装例

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

def solve(nums, queries):
    if not nums:
        return 0

    n = len(nums)
    # 隣接要素の差を格納する配列
    diff = [nums[i + 1] - nums[i] for i in range(n - 1)]

    # 連続する同じ差の長さを記録する配列
    rle = [0] * (n - 1)
    for i in range(n - 1):
        if i > 0 and diff[i] == diff[i - 1]:
            rle[i] = rle[i - 1] + 1
        else:
            rle[i] = 1

    ans = 0
    for i, j in queries:
        if i == j:
            ans += 1
        else:
            ans += rle[j - 1] >= (j - i)
    return ans

nums = [2, 4, 6, 8, 7, 6, 5, 2]
queries = [[3, 4], [0, 3], [2, 4]]
print(solve(nums, queries))

入力

[2, 4, 6, 8, 7, 6, 5, 2], [[3, 4], [0, 3], [2, 4]]

出力

2

計算量について

前処理（diff 配列と rle 配列の構築）には O(n) の時間がかかり、各クエリへの回答は O(1) で行えます。全体の計算量は O(n + q)（n は配列の長さ、q はクエリ数）となり、クエリごとに毎回 O(j - i) で判定する素朴な方法（最悪 O(n × q)）と比べて大幅に効率化されています。

まとめ

本記事では、隣接要素の差分とランレングス（連続長）の記録を組み合わせることで、等差数列判定クエリを高速に処理する手法を紹介しました。「区間クエリを前処理で高速化する」という定番テクニックの良い練習例なので、ぜひ自分でも実装してみてください。