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

数列の隣接の差（or 差の絶対値）の和、区間最大和を高速に求める

tadanori

数列$A_1, A_2, A_3, … , A_n$の隣接要素の差の和を高速に計算する方法です。AtCoderの問題などを解く場合に、数列の隣接要素の差（or差の絶対値）の総和を求める問題が結構出てきます。当然、愚直に解くと間に合わない要素数なので何か工夫しなければなりません。この記事では、高速化の考え方を含めて解説します。

Contents
  1. Ai – Ajの総和を求める
  2. |Ai – Aj|を求める（差の絶対値）
  3. おまけ
  4. まとめ

Ai – Ajの総和を求める

問題

問題

$A_1, A_2, … , A_n$という数列が与えられた場合に、以下の値を求めなさい

$$
\sum_{i=1}^{N} \sum_{j = i+1}^{N} (A_i – A_j)
$$

愚直解

愚直解は式通りの処理を実際に行うもので、以下になります。

このコードでは、iを0~N-1まで、jをi+1~N-1までループさせてa[i]-a[j]を計算し、結果に加算しています。

２重にループを行うので、計算量は（$O(N^2)$）になります。

n = 5
a = [5, 2, 4, 6, 1]

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

高速化

５つの値をa,b,c,d,eとして考えてみます。すると、$A_i – A_j$の計算を列挙すると以下のようになります。

$$\begin{eqnarray}
a-b,\ &a-c,\ &a-d,\ &a-e\\
&b-c,\ &b-d,\ &b-e\\
&&c-d,\ &c-e\\
&&&d-e\\
\end{eqnarray}$$

行にi, 列にjをとって整理したのが下図（右）です。

行で見るとaは4回、bは3回、cは2回、dは1回出現していることがわかります。また、列で見ると-bが1回、-cが2回、-dが3回、-eが4回出現していることがわかります。

これを整理すると、下図９（左）になります。

つまり、aはn-1回加算して0回減算、bはn-1-1回加算して1回減算することがわかります。

これをプログラムにすると以下のようになります。インデックスiは0スタートなことに注意してください。

２重ループ（$O(N^2)$）が$O(N)$となり高速化できました。

こんな感じで数式にして整理してみると、高速な解法がみつかったりします。


n = 5
a = [5, 2, 4, 6, 1]

print(a)

ans = 0
for i in range(n) :
    ans += a[i] * ((n-1-i) - i)

print(ans)

ランダムな数列で検算

実際に正しいかをランダムな値を10回生成して、愚直解と比較してチェックしてみます。

以下のプログラムは、愚直解と異なればNGを表示します。

import random

for _ in range(10) :
    n = random.randint(10, 100)
    a = [random.randint(1,100) for _ in range(n)]
    print("-"*100)
    print(len(a), a)

    ans1 = 0
    for i in range(n):
        for j in range(i+1, n):
            ans1 += a[i] - a[j]
        

    ans2 = 0
    for i in range(n) :
        ans2 += a[i] * ((n-1-i) - i)

    if ans1 == ans2 : print("OK")
    else: print("NG")

実行してみると以下のように全てOKになります。

----------------------------------------------------------------------------------------------------
24 [46, 74, 11, 2, 53, 87, 54, 2, 13, 19, 67, 61, 84, 87, 11, 64, 6, 68, 88, 28, 42, 44, 51, 53]
OK
----------------------------------------------------------------------------------------------------
15 [81, 72, 6, 63, 80, 67, 81, 84, 54, 57, 48, 3, 86, 1, 57]
OK
----------------------------------------------------------------------------------------------------
25 [69, 81, 26, 65, 43, 69, 86, 32, 73, 65, 60, 34, 19, 61, 21, 49, 78, 24, 57, 80, 95, 51, 80, 17, 33]
OK
----------------------------------------------------------------------------------------------------
50 [78, 37, 65, 65, 20, 30, 76, 16, 81, 23, 6, 91, 50, 41, 57, 10, 74, 86, 74, 2, 13, 88, 85, 13, 46, 15, 16, 9, 2, 69, 40, 77, 38, 75, 44, 81, 44, 48, 95, 32, 86, 9, 76, 34, 97, 77, 87, 81, 79, 78]
OK
----------------------------------------------------------------------------------------------------
83 [64, 91, 20, 20, 59, 25, 16, 2, 25, 93, 28, 19, 81, 41, 5, 75, 81, 87, 8, 21, 75, 91, 68, 38, 42, 16, 31, 39, 31, 31, 79, 90, 22, 29, 54, 74, 68, 99, 40, 23, 20, 89, 28, 57, 64, 61, 68, 91, 4, 59, 10, 31, 57, 6, 33, 43, 15, 92, 67, 47, 78, 26, 86, 55, 16, 18, 36, 48, 49, 9, 59, 59, 80, 4, 96, 90, 84, 66, 67, 18, 23, 97, 17]
OK
----------------------------------------------------------------------------------------------------
42 [89, 17, 78, 32, 10, 7, 11, 60, 1, 46, 52, 3, 76, 94, 12, 33, 87, 86, 74, 83, 30, 46, 21, 64, 91, 95, 96, 76, 52, 71, 93, 87, 100, 74, 1, 24, 8, 32, 97, 75, 71, 13]
OK
----------------------------------------------------------------------------------------------------
59 [27, 6, 95, 84, 39, 2, 61, 87, 75, 58, 84, 90, 22, 93, 78, 85, 79, 88, 78, 31, 63, 44, 3, 25, 44, 2, 31, 65, 43, 1, 100, 66, 48, 2, 92, 40, 23, 69, 8, 64, 13, 41, 1, 52, 26, 4, 4, 13, 17, 32, 86, 82, 33, 13, 98, 56, 36, 75, 18]
OK
----------------------------------------------------------------------------------------------------
36 [26, 70, 100, 17, 62, 46, 13, 61, 59, 50, 58, 4, 96, 41, 46, 35, 83, 91, 39, 5, 99, 9, 72, 79, 24, 10, 77, 86, 86, 66, 43, 55, 23, 85, 88, 40]
OK
----------------------------------------------------------------------------------------------------
30 [11, 19, 77, 44, 46, 93, 37, 34, 41, 22, 79, 72, 68, 97, 31, 39, 28, 18, 83, 77, 79, 43, 34, 75, 87, 97, 67, 12, 56, 49]
OK
----------------------------------------------------------------------------------------------------
36 [65, 54, 96, 56, 68, 23, 90, 93, 3, 80, 100, 16, 10, 48, 11, 95, 10, 77, 34, 30, 64, 86, 100, 28, 49, 93, 64, 71, 86, 63, 49, 100, 72, 40, 29, 36]
OK

|Ai – Aj|を求める（差の絶対値）

問題

問題

$A_1, A_2, … , A_n$という数列が与えられた場合に、以下の値を求めなさい

$$
\sum_{i=1}^{N} \sum_{j = i+1}^{N} |A_i – A_j|
$$

愚直解

愚直回は、先ほどと同じでabs(a[i] - a[j])を求めるものになります。

n = 5
a = [5, 2, 4, 6, 1]

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

高速化

高速化の仕方は同じです。異なるのは、配列aをあらかじめソートしている点です。

こうすることでa[i]-a[j]は必ず正になります。

このように、あらかじめソートしておくことで絶対値を省くことができます。


n = 5
a = [5, 2, 4, 6, 1]

a.sort(reverse=True)

ans = 0
for i in range(n) :
    ans += a[i] * ((n-1-i) - i)

print(ans)

ランダムな数列で検算

実際に正しいかをランダムな値を10回生成して、愚直解と比較してチェックしてみます。

以下のプログラムは、愚直解と異なればNGを表示します。

本当にソートしても良いかも確認するために、愚直回を計算した後にsortしている点に注意してください。

import random

for _ in range(10) :
    n = random.randint(10, 100)
    a = [random.randint(1,100) for _ in range(n)]
    print("-"*100)
    print(len(a), a)

    ans1 = 0
    for i in range(n):
        for j in range(i+1, n):
            ans1 += abs(a[i] - a[j])
        

    a.sort(reverse=True)
    ans2 = 0
    for i in range(n) :
        ans2 += a[i] * ((n-1-i) - i)

    if ans1 == ans2 : print("OK")
    else: print("NG")

こちらも問題なく全てOKになりました。

----------------------------------------------------------------------------------------------------
50 [47, 33, 63, 7, 66, 94, 92, 1, 57, 80, 22, 22, 71, 41, 55, 13, 10, 74, 15, 34, 79, 54, 78, 52, 23, 76, 79, 54, 40, 61, 3, 89, 8, 86, 73, 91, 66, 78, 39, 7, 85, 48, 72, 62, 29, 21, 89, 69, 10, 33]
OK
----------------------------------------------------------------------------------------------------
67 [13, 5, 61, 76, 100, 40, 55, 77, 34, 4, 1, 72, 28, 41, 79, 83, 56, 98, 74, 58, 72, 77, 57, 24, 68, 57, 84, 57, 76, 2, 23, 53, 83, 21, 5, 17, 98, 8, 80, 23, 82, 59, 65, 67, 6, 57, 16, 80, 80, 9, 19, 94, 28, 54, 47, 46, 17, 100, 13, 17, 74, 38, 43, 69, 26, 33, 92]
OK
----------------------------------------------------------------------------------------------------
69 [76, 92, 84, 52, 35, 75, 42, 47, 9, 89, 45, 93, 56, 99, 11, 10, 96, 13, 16, 47, 40, 57, 22, 72, 3, 70, 32, 21, 16, 77, 7, 60, 24, 20, 50, 55, 39, 65, 62, 77, 55, 41, 97, 61, 59, 59, 40, 99, 25, 37, 19, 65, 4, 99, 64, 13, 43, 52, 21, 9, 47, 43, 32, 81, 51, 55, 42, 39, 90]
OK
----------------------------------------------------------------------------------------------------
48 [37, 9, 22, 6, 99, 77, 23, 57, 15, 26, 51, 88, 32, 77, 12, 30, 27, 66, 20, 30, 63, 11, 8, 93, 54, 87, 78, 54, 79, 81, 46, 26, 21, 12, 90, 72, 1, 88, 61, 27, 72, 75, 50, 48, 64, 36, 24, 36]
OK
----------------------------------------------------------------------------------------------------
55 [28, 33, 2, 1, 43, 45, 77, 90, 92, 6, 39, 21, 4, 9, 82, 82, 26, 60, 71, 83, 65, 98, 55, 18, 88, 39, 61, 73, 78, 2, 44, 80, 27, 24, 26, 84, 30, 24, 33, 100, 23, 80, 73, 41, 96, 61, 27, 65, 27, 93, 33, 1, 24, 21, 22]
OK
----------------------------------------------------------------------------------------------------
76 [21, 48, 8, 48, 65, 27, 18, 41, 76, 57, 85, 29, 69, 20, 14, 15, 29, 75, 5, 70, 31, 74, 3, 14, 24, 17, 43, 49, 25, 88, 60, 57, 25, 70, 90, 96, 25, 56, 9, 36, 46, 35, 35, 10, 31, 97, 53, 92, 43, 37, 13, 36, 11, 1, 18, 44, 35, 4, 84, 21, 43, 32, 31, 81, 1, 76, 28, 18, 80, 13, 68, 74, 15, 87, 10, 79]
OK
----------------------------------------------------------------------------------------------------
94 [92, 68, 85, 99, 70, 29, 99, 89, 22, 88, 46, 15, 66, 31, 63, 35, 88, 24, 98, 72, 12, 20, 17, 80, 61, 55, 43, 26, 31, 39, 73, 43, 56, 55, 97, 2, 69, 20, 74, 71, 84, 10, 30, 1, 16, 100, 22, 65, 1, 28, 54, 3, 33, 71, 94, 82, 59, 72, 52, 42, 56, 48, 64, 74, 29, 57, 43, 61, 58, 66, 46, 1, 39, 40, 61, 76, 77, 40, 21, 100, 58, 65, 20, 96, 42, 42, 95, 40, 72, 68, 49, 27, 14, 53]
OK
----------------------------------------------------------------------------------------------------
95 [32, 21, 83, 50, 79, 36, 8, 57, 84, 80, 37, 54, 7, 25, 51, 25, 76, 68, 40, 20, 43, 20, 13, 72, 98, 61, 36, 71, 34, 1, 92, 27, 76, 19, 45, 21, 80, 77, 98, 38, 30, 30, 50, 45, 35, 94, 19, 82, 44, 26, 28, 3, 53, 30, 42, 48, 99, 17, 95, 67, 44, 48, 78, 99, 97, 76, 61, 26, 1, 78, 30, 53, 58, 88, 65, 23, 73, 38, 29, 1, 43, 23, 70, 41, 29, 39, 63, 56, 46, 85, 33, 81, 95, 36, 11]
OK
----------------------------------------------------------------------------------------------------
59 [72, 10, 81, 90, 43, 44, 7, 79, 17, 47, 79, 9, 13, 87, 8, 71, 28, 56, 32, 49, 61, 20, 53, 15, 75, 50, 55, 16, 55, 54, 82, 58, 18, 35, 82, 55, 47, 35, 52, 47, 96, 60, 49, 16, 25, 57, 79, 16, 65, 52, 37, 68, 68, 100, 59, 88, 51, 95, 72]
OK
----------------------------------------------------------------------------------------------------
43 [65, 76, 27, 31, 24, 42, 81, 100, 32, 46, 43, 66, 20, 34, 80, 10, 26, 6, 87, 88, 86, 8, 98, 62, 26, 45, 85, 74, 19, 61, 1, 72, 19, 87, 20, 64, 66, 57, 55, 20, 20, 92, 62]
OK

おまけ

積の総和（ΣΣAi×Aj）

問題

$A_1, A_2, … , A_n$という数列が与えられた場合に、以下の値を求めなさい

$$
\sum_{i=1}^{N} \sum_{j = i+1}^{N} A_i \times A_j
$$

解法

$A_1$に対する乗算、$A_2$に対する乗算と見ていくと以下のようにまとめることができます。

$$\begin{eqnarray}
A_1\times A_2 + A_1\times A_3 + A_1\times A_4 + … + A_1 \times A_n\\
= A_1(A_2 + A_3 + A_4 + … + A_n)\\
A_2\times A_3 + … + A_2 \times A_n\\
= A_2(A_3 + ….. + A_n)\\
\end{eqnarray}$$

つまり、$A_1$に対しては$\sum_{i = 2}^{n}$までの総和と乗算、$A_2$については、$\sum_{i = 3}^{n}$までの乗算となります。

あらかじめ$A_1$ ~ $A_n$までの総和を求めていれば、引きながら掛け算をしていくだけで求めることができます。

n = 5
a = [5, 2, 4, 6, 1]

tot = sum(a)
ans = 0
for i in range(n) :
    tot -= a[i]
    ans += a[i] * tot

区間和の最大値を求める

$A_1, A_2, …, A_n$の区間[l, r]の和の最大値を求める問題です。この問題は累積和を計算して考えると分かりやすいです。

数列A = [1, -2, -3, 4, 5, -10, 3, 7, -2, 9]の場合について考えます。累積和は下図のようになります。これを図示して見ると、区間和の最大は区間差が一番大きい部分になることがわかります。

i番目までみた場合の累積和の最小値から最大値までが最大です

区間和の最大のイメージ図

プログラム

愚直解と、高速化したものです。

愚直にやると、開始位置をiにして最後まで加算しながら最大値を求める処理を行うことになるため、２重ループになります。

一方、高速な手法は、O(N)で計算できています。

a = [1, -2, -3, 4, 5, -10, 3, 7, -2, 9]

# 愚直解
ans = sum(a)
for i in range(len(a)):
    tot = 0
    for j in range(i, len(a)) :
        tot += a[j]
        ans = max(ans, tot)
    
print(ans)


# 高速な解法
a = [1, -2, -3, 4, 5, -10, 3, 7, -2, 9]
cur_sum, min_sum = 0,0
ans = sum(a)
for i in range(len(a)) :
    cur_sum += a[i]
    min_sum = min(min_sum, cur_sum)
    ans = max(ans, cur_sum-min_sum)

print(ans)

まとめ

数列の隣接差の総和を求める問題、考え方さえ分かれば簡単ですが、たまに忘れてちょっと「？」となってしまうので、まとめてみました。

メールアドレスが公開されることはありません。 ※ が付いている欄は必須項目です

Recommend
こちらの記事もどうぞ
記事URLをコピーしました