計算量を削減して a+b+c=m の組み合わせを数える方法｜O(N³)→O(N)まで段階的に解説【Python】

Aru

競技プログラミング（AtCoderなど）で頻出する、計算量を意識したアルゴリズムの基礎を解説します。今回は、a + b + c = m の組み合わせ数を求める問題を題材に、$O(N^3)$の全探索から$O(N)$まで計算量を削減する方法を段階的に見ていきます。変数削減や範囲計算といった基本テクニックを理解するための入門的な内容です。

Contents

問題設定
愚直解（3重ループ）
計算量の削減（2重ループへ）
さらに計算量を削減（単一ループへ）
まとめ

問題設定

今回の問題は、$1 \le a, b, c \le N$ の範囲が与えられた場合に、$a + b + c = m$ を満たす整数の組 $(a, b, c)$ がいくつあるかを求めるというものです。

問題

$1 \le a, b, c \le N$ の範囲が与えられた場合に、$a + b + c = m$ を満たす整数の組 $(a, b, c)$ がいくつあるかを求めてください。

条件はシンプルですが、与えられる $N$ の大きさによって、適切な実装方法（計算量）を選択する必要があります。

今回の問題は、単なる数え上げに見えますが、競技プログラミングでは非常に重要な「計算量削減の典型パターン」を含んでいます。特に、

複数変数の組み合わせをそのまま全探索すると計算量が爆発する
条件式を変形することで、探索する変数の数を減らせる
さらに、数え上げを「ループ」ではなく「範囲の長さの計算」に置き換えられる

といった考え方は、多くの問題に応用可能です。

また、この問題はシンプルな形をしているため、これらのテクニックを段階的に確認するのに適した題材です。

リンク

愚直解（3重ループ）

まずは一番わかりやすい素直な解法です。プログラミング初心者がまず思いつくものだと思います。

個人的には、まずは愚直解を書いてみるのは良いと思います。「愚直解なら分かる」場合も結構多いし、何より、高速化したコードの検算に使えます。愚直解を使って、よくわからないけどWAになる場合のバグを発見することができます。

$N$ が小さい場合（例えば $N \le 100$ 程度）であれば、3重ループですべての組み合わせを探索する素直な実装で問題ありません。

# O(N^3)の解法
count = 0
for a in range(1, N + 1):
    for b in range(1, N + 1):
        for c in range(1, N + 1):
            if a + b + c == m:
                count += 1

直感的に理解しやすいコードですが、計算量は $O(N^3)$ となります。

計算量の削減（2重ループへ）

$N$ が大きくなると3重ループでは処理時間が長くなります。

例えば $N=1000$ の場合、ループ回数は $10^9$ 回となり、Python環境や競技プログラミングの実行時間制限（TLE）に引っかかってしまいます。

そこで、先ほどの式 $a + b + c = m$ を以下のように変形してみます。

$$c = m – (a + b)$$

このように書き直せば、$a$ と $b$ の値を決めれば、$c$ の値は一意に計算できることがわかります。これを利用してループを1つ減らしてみます。

ここで注意するのは、「算出した$c$が条件範囲の値かどうか」のチェックが必要だということです。例えば、$a, b = 100, 100$で$m=100$の場合、$c = 100-(100+100)=-100$となるので、これは$1\le c \le N$の範囲ではありませんのでカウントしないようにしなければなりません。

# O(N^2)の解法
count = 0
for a in range(1, N + 1):
    for b in range(1, N + 1):
        c = m - (a + b)
        
        # 算出した c が 1 から N の範囲に収まっているかを確認
        if 1 <= c <= N:
            count += 1

これにより、$c$ を探索するループを消すことができ、計算量を $O(N^2)$ まで削減できました。$N=1000$の場合でもループは $10^6$ 回で済むため、計算時間が$1/1000$になります。

さらに計算量を削減（単一ループへ）

先ほどのやり方でも、$N=10^5$ などになると $O(N^2) = 10^{10}$ となり、やはり処理時間が膨大になります。この場合は、さらにループを減らして $O(N)$ に落とし込む方法を考えます。

ここでは、少し数学的に変数の「範囲」を絞り込んでみます。

$a$ を固定したとします。このとき、残りの変数について $b + c = m – a$ となります。ここで見通しを良くするために $X = m – a$ と置くと、問題は「$b + c = X$ となる組み合わせはいくつあるか？」という形に分解できます。

さらに $c = X – b$ と変形し、$c$ が取り得る条件（$1 \le c \le N$）に当てはめます。

$$1 \le X – b \le N$$

これを $b$ について整理すると、以下のようになります（符号に注意して展開します）。

$$X – N \le b \le X – 1$$

同時に、もともとの $b$ の条件である $1 \le b \le N$ も満たす必要があります。つまり、$b$ が存在できる値の範囲は、これら2つの条件の共通範囲となります。

下限: $1$ と $X – N$ の大きい方 $\to \max(1, X – N)$
上限: $N$ と $X – 1$ の小さい方 $\to \min(N, X – 1)$

したがって、ある $X$ に対して条件を満たす $b$ の個数は、「上限から下限を引いて1を足した数」になります（※上限と下限が逆転する場合は条件を満たす $b$ が存在しないため 0 個になります）。

この計算はループを使わずに $O(1)$ で求めることができます。あとは $a$ をループさせるだけです。

# O(N)の解法
count = 0
for a in range(1, N + 1):
    X = m - a
    
    # bの下限と上限を計算
    min_b = max(1, X - N)
    max_b = min(N, X - 1)
    
    # 条件を満たすbの個数を足す（存在しない場合は0）
    if min_b <= max_b:
        count += (max_b - min_b + 1)

この工夫により、計算量を $O(N)$ まで落とすことができました。これなら $N=10^7$ のような巨大な数が与えられても、問題なく処理が可能です。