n 個の数列から異なる 2項を選んで作る積の総和

　LeetCode の 「2316. Count Unreachable Pairs of Nodes in an Undirected Graph」で、「n 個の数列から異なる 2項を選んで作る積の総和」を求める必要がありました。

　2重ループをまわせば簡単に実装できますが、n = 10^5 になると計算量が O(10^10) となりTLE（制限時間オーバー）となります。

そこで今回は、TLEにならない方法を python3 で以下に実装してみます。 

n = 4
A = [4, 2, 1, 3]

ans = 0
firstCount = A[0] # 初項の値をfirstCountに代入
for i in range(1, n): # 1項目から始める
    ans += firstCount * A[i] # firstCountに現在みている項の値を掛けたものを答えに足す
    firstCount += A[i]  # firstCount に現在みている項の値を足して更新する
print(ans) #--> 35

上のコードで中で示した数列は、A でその要素数( n ) は 3です。
この数列から「異なる2項を選んで作る積の総和」とは、
(4 * 2 + 4 * 1 + 4 * 3) + (2 * 1 + 2 * 3) + (1 * 3) = 35 ・・・（式1）
ということです。

この式は、まさに 2重ループであり、実装してみると、

ans = 0
for i in range(n-1):
    for j in range(i+1,n):
        ans += A[i] * A[j]
print(ans)

となります。
しかし、先に述べたようにこれでは n = 10^5 になると計算量が O(10^10) となりTLE（制限時間オーバー）となります。

よって、（式1）を変形して、効率的に計算できる方法を以下に考えてみます。
結論から言えば（式1）は以下のように変形できます。

4 * 2 + (4 + 2) * 1 + (4 + 2 + 1) * 3 = 35・・・（式2）

数列を {a, b, c, d}　として、一般化してみると、（式2）の左辺に当たる部分は、

(a * b) + (a + b) * c + (a + b + c) * d

と、規則性を持った式になっています。

よって、最初のコードのように 1重の ループで実装できるようになります。

これで計算量は、O(n)程度となり、n = 10^5 は大丈夫になります。


お役に立てれば幸いです。