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 は大丈夫になります。
お役に立てれば幸いです。