18 数列の和 Dif:386 ABC177 C：「AtCoder 凡人が『緑』になるための精選50問詳細解説」サンプル

この記事は拙著「AtCoder　凡人が『緑』になるための精選50問詳細解説」のサンプルです
価格：100円
kindle：https://www.amazon.co.jp/dp/B09C3TPQYV
booth(pdf)：https://sano192.booth.pm/items/3179185

前：https://qiita.com/sano192/items/d6d2397f5f1bde156631
次：https://qiita.com/sano192/items/b301698295e6b6721b35

【目標】
・数列の和を求める問題は具体例で考えるということを覚える
・計算量を減らせるように数式を変形する方法を身につける

【概要】
数列が与えられ、その組に対してなにかしら処理したものの和を出すというのは頻出の問題である。本問のように愚直に全ての組み合わせを計算すると計算量が大きすぎてTLE(実行時間制限超過)する、というのもお決まりのパターンだ。そしてこの手の問題は解説を見てもシグマをやたら並び替えたり変換したりしていて何をしているのか、まったくわからない・・・。
この問題で数列の和問題はどのように解くのかを学ぼう。

【方針】
数列が与えられ、その組に対してなにかしら処理したものの和を出す、という問題が出てきた場合、まず最初に入力例を見る。たいてい入力例1には具体的に何を計算するのか丁寧に書いてくれているからだ。

次に紙とペンを用意し、入力例1に書いている計算の内容をそのまま紙に写す。

ここまでを必ず行うこと。AiだのAjだののまま、頭の中でΣをこねくり回して解こうとしても絶対に解けない。
とにかく具体例、そして紙とペンだ。

ではこの問題の入力例1をみよう。
「入力例1」
3
1 2 3

N=3
A：1 2 3

計算方法も丁寧に記載がある。

まず紙に計算式をそのまま写す。
12+13+2*3

紙に写したら次に考えるのは「どこをかっこでくくれるか」。
最初の2項について1(=A1)が重複しているからこれでくくれそうだ。
1*(2+3)+2*3
(2+3)=A2+A3だから、ここをうまくやれれば計算量を減らせそう。しかしまだよくわからない。

入力例を見てもよくわからない、というときは自分でもっと大きなサンプルを作る。
A：1 2 3 4
ならどうなるか。

12+13+14+23+24+34
1、2を使ってかっこでくくろう。
=1*(2+3+4)+2*(3+4)+34
(2+3+4)=A2+A3+A4
(3+4)=A3+A4
となっている。どうやら
A1(A2~ANの和)+A2*(A3～ANの和)+...
と計算できそうだ。
A2~ANの和
A3~ANの和
...
はそれぞれ
A1～ANの和-A1
A2～ANの和-A2
…
となるからこれなら逐一sum関数で和を計算せずに済む。この方針でTLE(実行時間制限超過)せずに解ける。

この変形が一般の場合において正しいかどうかの証明はコンテスト中にしなくてもよい。とりあえずこの方針でコードを作ってみて入力例を試し、出力が合っていれば提出する。
証明はコンテストの後にゆっくり考えよう。考えてわからなくても無理に証明まで理解する必要はない。

【実装】
入力を受け取り、ついでにmod=10^9+7を定義しておく。

N=int(input())
A=list(map(int, input().split()))
mod=10**9+7

次に
A1*(A2~ANの和)+A2*(A3～ANの和)+...
すなわち
A[0](A[1]~A[N-1]の和)+A[1](A[2]~A[N-1]の和)+...
を計算する。(リストAは0インデックスなのでA1→A[0]、A2→A[1]、...とずれる)

(A[1]~A[N-1]の和)=sum(A)-A[0]　
(A[2]~A[N-1]の和)=(A[1]~A[N-1]の和)-A[1]　
(A[3]~A[N-1]の和)=(A[2]~A[N-1]の和)-A[2]　
…
であるからまずAの合計を計算し、A[1]、A[2]、...を順に引いていけば(A[1]~A[N-1]の和)、(A[2]~A[N-1]の和)、...をO(1)で求めることができる。
足し算しながら10^9+7で余りを取ることを忘れずに。

A_sum=sum(A)

ans=0

for i in range(N):
    A_sum-=A[i]
    ans+=A[i]*(A_sum)
    ans%=mod

最後に答えを出力して終わり。

print(ans)

【コード全文】

N=int(input())
A=list(map(int, input().split()))
mod=10**9+7

A_sum=sum(A)

ans=0

for i in range(N):
    A_sum-=A[i]
    ans+=A[i]*(A_sum)
    ans%=mod

print(ans)

この記事は拙著「AtCoder　凡人が『緑』になるための精選50問詳細解説」のサンプルです
価格：100円
kindle：https://www.amazon.co.jp/dp/B09C3TPQYV
booth(pdf)：https://sano192.booth.pm/items/3179185

前：https://qiita.com/sano192/items/d6d2397f5f1bde156631
次：https://qiita.com/sano192/items/b301698295e6b6721b35