2024.06.11 2024.08.28

記事内に商品プロモーションを含む場合があります

数列問題｜要素2つの合計が一定値以上になる組み合わせを高速に求める方法

Aru

この記事では、A1, A2, A3…Anという数列から選んだ２つの要素を足してn以上になるパターンの数を高速にカウントする方法について解説します。具体例には、Pythonを使って、この問題を高速に解くアルゴリズムを実装する手順を紹介します。また、愚直解を与えて、生成AIで高速化する試みも行ってみました。

Contents

対象問題
愚直解
高速化した解法
まとめ

対象問題

対象とする問題は以下になります。

n個の数列$A_1, A_2, A_3…A_n$の2つの要素を足して100以上になる組み合わせの数はいくつありますか？

今回は「100以上」にしましたが、ここは「n以上」と一般化しても解法は同じです。

愚直解

まずは、愚直解です。愚直解は、２重ループを使って以下のように記述可能です。

下記のプログラムでは、要素数n=10とし、1から100の乱数を発生させて数列aを生成し、100以上かどうかを判定し100以上であればcntをインクリメントします。

import random

n = 10
a = [random.randint(1, 100) for i in range(n)]

print(a)
cnt = 0
for i in range(n):
    for j in range(i+1, n):
        print(a[i], a[j], end=" ")
        if a[i] + a[j] >= 100 :
            cnt+=1
            print(">=100", "cnt=", cnt)
        else :
            print("")
print(cnt)

プログラムでは途中経過も同時に出力するようにしていますが、結果は例えば以下のようになります。

[14, 25, 54, 21, 90, 21, 15, 17, 39, 75]
14 25 
14 54 
14 21 
14 90 >=100 cnt= 1
14 21 
14 15 
14 17 
14 39 
14 75 
25 54 
25 21 
25 90 >=100 cnt= 2
25 21 
25 15 
25 17 
25 39 
25 75 >=100 cnt= 3
54 21 
54 90 >=100 cnt= 4
54 21 
54 15 
54 17 
54 39 
54 75 >=100 cnt= 5
21 90 >=100 cnt= 6
21 21 
21 15 
21 17 
21 39 
21 75 
90 21 >=100 cnt= 7
90 15 >=100 cnt= 8
90 17 >=100 cnt= 9
90 39 >=100 cnt= 10
90 75 >=100 cnt= 11
21 15 
21 17 
21 39 
21 75 
15 17 
15 39 
15 75 
17 39 
17 75 
39 75 >=100 cnt= 12
12

上記の数列[14, 25, 54, 21, 90, 21, 15, 17, 39, 75]の場合は、100を超える組み合わせは12個になります。

愚直解は２重ループになっているので、計算量は$O(N^2)$となります。

高速化した解法

解法①

高速化するには、ループ数を減らす必要があります。

数列の２つの組み合わせだけに注目した場合、数列をソートしても結果は変わりません（組み合わせパターンの数は変化しません）。

なので、ソートして考えます。

現在注目するAを$A_i$、比較するAを$A_r$とします。$r$の初期値は一番大きな値にセットします。次に、A[i] + A[r]が100未満になるまでrを小さい方に移動させます。

するとN-(r+1)が、$A_i$と加算すると100を超える数になります。

これをiが0からn-1まで繰り返せば、個数がわかります。

なお、rは初期化せずにそのまま続けることができます（i+1の要素はiより常に大きいのでrより右側とは足して100になる）。

注意点としてはrがiより小さくならないようにすることです（プログラム中のmax(i, r)）。

これで計算量$O(n)$で個数を数えることができます。

ただ、実際には、事前にソートが必要なので、計算量は全体では$O(nlogn)$になります。

import random

n = 10
a = [random.randint(1, 100) for i in range(n)]
a.sort()

r = n-1
tot = 0
for i in range(n) :
    while (r >= 0 and a[i] + a[r] >= 100) :
        r -= 1
    r = max(i, r)
    tot += n - (r + 1)

print(tot)

解法②

考え方は同じで、「ソートして、足して100以上になる要素の場所を調べ、右側にある要素数をカウント」します。

違うのは、要素の数を調べる部分です。

こちらの方法では、100以上となる境界を二分探索で求めます（bisectを利用）。

import bisect
import random

n = 10
a = [random.randint(1, 100) for i in range(n)]
a.sort()

tot = 0
for i in range(n):
    target = 100 - a[i]
    pos = bisect.bisect_left(a, target, i + 1)
    tot += n - pos

print(tot)

この解法の計算量は$O(nlogn)$になります。

最近は、このレベルの問題なら指示が適切ならChatGPTで解けるようになってきました。「二分探索で求めてください」などの指示が必要なので、解き方を知っておいた方が期待する解答に近づけます。

この計算を応用して解けるとして、

AtCoder Beginner Contest 353のC問題

があります。興味がある方は解いてみてください。

番外編（生成AIを使って最適化）

他の記事にも書きましたが私はCodeiumという生成AIをコーディングのアシストとして使っています。

あわせて読みたい

無料で使えるコード生成AI「Codeium」をVSCodeで使う方法

このCodeiumには、Refactorという機能があり、その中に「Make faster and more efficient（より速く、より効率的に変換）」というメニューがあります（下図）。

Refactorのメニュー — Refactorをクリックした時に表示されるメニュー

これを使うと、変換されたコードが出力されます（実際には、差分が表示されます。ここでacceptをクリックすると修正したコードが反映されます）。

この機能を使って、以下の愚直解を最適化してみました。

def count(a) :
    n = len(a)
    cnt = 0
    for i in range(n):
        for j in range(i+1, n):
            if a[i] + a[j] >= 100 :
                cnt+=1
    return cnt

結果は以下になります。このコード、解法①とほぼ同じです。この程度の最適化は、生成AIのアシスト機能で自動で作れるようになったみたいです。生成AI恐るべしです。

def count(a) :
    a.sort()
    l, r = 0, len(a)-1
    cnt = 0
    while l < r:
        if a[l] + a[r] >= 100:
            cnt += r-l
            r -= 1
        else:
            l += 1
    return cnt

この例を考えると、とりあえず愚直解を書いてみて、生成AIで最適化させてみるというのこれからの１つの手段かもしれません。

とはいえ、出力されたコードを理解するだけのプログラミングとアルゴリズムスキルは必要だと思います。生成AIがあるからといって、完全に任せるのは難しいです。