AtCoder Beginner Contest 163 D を Python で

AtCoder Beginner Contest 163 D を Python で解説していきたいと思います。
結構数学的な問題なのかな？と思ったのですがどうだったでしょうか。
それでは，まずは問題文から。
https://atcoder.jp/contests/abc163/tasks/abc163_d

D – Sum of Large Numbers

実行時間制限: 2 sec / メモリ制限: 1024 MB

配点 : 400 点

問題文

10^100, 10^100+1, …, 10^100+N の N+1 個の数があります。

この中から K 個以上の数を選ぶとき、その和としてあり得るものの個数を mod(10^9+7) で求めてください。

制約

1≤N≤2×10^5
1≤K≤N+1
入力は全て整数

入力

入力は以下の形式で標準入力から与えられる。

N K

出力

和としてあり得るものの個数を mod(10^9+7)で出力せよ。

入力例 1

3 2

出力例 1

以下の 10 通りが考えられます。

(10^100)+(10^100+1)=2×10^100+1
(10^100)+(10^100+2)=2×10^100+2
(10^100)+(10^100+3)=(10^100+1)+(10^100+2)=2×10^100+3
(10^100+1)+(10^100+3)=2×10^100+4
(10^100+2)+(10^100+3)=2×10^100+5
(10^100)+(10^100+1)+(10^100+2)=3×10^100+3
(10^100)+(10^100+1)+(10^100+3)=3×10^100+4
(10^100)+(10^100+2)+(10^100+3)=3×10^100+5
(10^100+1)+(10^100+2)+(10^100+3)=3×10^100+6
(10^100)+(10^100+1)+(10^100+2)+(10^100+3)=4×10^100+6

入力例 2

200000 200001

出力例 2

全てを選ぶしかないので 1 通りです。

入力例 3

141421 35623

出力例 3

220280457

解説

10^100, 10^100+1, …, 10^100+N の N+1 個の数があり，この中から K 個以上の数を選ぶとき、その和としてあり得るものの個数を mod(10^9+7) で求めるというものです。

ただし，ここで 10^100 は見掛け倒しです。（笑）
ただの和を考えればいいので，例えば 10^100 + 10^100+1 = 2 × 10^100 + 1 は，結局１に対応します。

つまり，０～N の数があり，この中から K 個以上の数を選ぶときの和のありうるものの個数に帰着されます。ただし，K は 10^100 の個数に対応するので，
（先ほどの例では10^100 + 10^100+1 = 2 × 10^100 + 1 = K × 10^100 + 1)
ここでは０～N の数の和自体が同じでも， K の値が異なれば異なる和と考えます。
以上より，以下が導けます。

K ≦ i ≦ N+1 における i において，０～N の数から i 個の数を選ぶときの和のありうるものの個数を，すべての i に対して足し合わせたものを出力する。

ここで問題となるのは，各 i における i 個の数を選ぶときの和のありうるものの個数をどう出すかですが，この０～Nの値の和は，i 個の数の和の最小値から最大値まで全ての値を取ることができます。
これは数学的に証明できますが，ここでは省略します。
よって，このときの和の個数は，（最大値　ー　最小値　＋１）で出すことができます。

最大値は，N から i 個取った数の和ですし，最小値は０から i 個取った数の和です。
ここで，１～N までの数の和は，N(N+1)/2 で与えられるので，これを使います。

まず最小値は，０から i 個取った数の和，すなわち０から i-1 までの数の和なので，上の式から
min = (i-1)×i/ 2
で得られます。
次に最大値ですが，これは N から i 個取った数の和，すなわち N-i+1 から N までの和になります。色々な計算の仕方がありますが，一番やりやすいのが，1からN までの和から，１からN-i までの和を引けばいいです。
max = N×(N+1)/2 – (N – i)×(N – i + 1)/2
になります。よって，あとは以上から得られた max – min + 1 を，K ≦ i ≦ N+1 なる全ての i において足し合わせた結果を出力してやります。これでO(n)で解くことができました。
それでは以下に解答例を載せておきます。
mod(10^9+7)で出力するので，最後の結果を 10^9+7 で割った余りを出力することに注意してください。

解答例

n,k = map(int, input().split())
 
ans = 0
mod = 10**9+7
 
for i in range(k,n+2):
    ans += int((n*(n+1)/2 - (n-i)*(n-i+1)/2 - (i-1)*i/2 +1))
 
print(ans%mod)

いかがだったでしょうか。
コード自体はかなり簡潔に書けますが，考え方がかなり数学チックでややこしい問題でしたね。
もし分からないところがあればコメントしてくださいね！
他の問題も解説しているので，ぜひご覧ください。
それでは！

追記

０～Nの値の和は，i 個の数の和の最小値から最大値まで全ての値を取ることができます。
これは数学的に証明できますが，ここでは省略します。

上記のこの証明の部分を知りたいと言う方がおられたので，追記しておきたいと思います。
i の値を固定して考えます。i は K ≦ i ≦ N+1，K は1≤K≤N+1 なる数です。
まず上記の説明からmin と max が与えられます。
min + 1 の値をとれるかについてですが，i は高々 N+1 であり，０からN までの N+1 個の数の値について考えているので，i = N+1 なら min と max が一致します。
i が N 以下である時は，０からN までで選んだ i 個の数のうち，その数のうちで最大となるもの(m_1とする)より１大きい数(m_2とする)が必ず選べるので，m_1とm_2を入れ替えることでとmin+1 の値は取れます。
次の min + 2 は，先ほどの m_1 – 1 なる数字と m_2 を入れ替えることで作れます。
以下，帰納的に，min + j = max となるような j が得られるまで，min + 2 , min + 3 … min + j の値は得られることが分かります。
以上より，厳密ではありませんが，０～Nの値の和は，i 個の数の和の最小値から最大値まで全ての値を取ることができることが帰納的に示すことができました。

タムログ

プログラミングについての情報発信をしていきます！