ABC163 D - Sum of Large Numbers

備忘録

問題

atcoder.jp

回答

import sys
import os
import math
import string

ii = lambda: int(sys.stdin.buffer.readline().rstrip())
il = lambda: list(map(int, sys.stdin.buffer.readline().split()))
fl = lambda: list(map(float, sys.stdin.buffer.readline().split()))
iln = lambda n: [int(sys.stdin.buffer.readline().rstrip()) for _ in range(n)]

iss = lambda: sys.stdin.buffer.readline().decode().rstrip()
sl = lambda: list(map(str, sys.stdin.buffer.readline().decode().split()))
isn = lambda n: [sys.stdin.buffer.readline().decode().rstrip() for _ in range(n)]

lcm = lambda x, y: x * y / math.gcd(x, y)

MOD = 10 ** 9 + 7


def main():
    if os.getenv("LOCAL"):
        sys.stdin = open("input.txt", "r")

    N, K = il()
    ret = 0
    for k in range(K, N + 2):
        ret += (2 * N - k + 1) * k // 2 - (k - 1) * k // 2 + 1
    print(ret % MOD)


if __name__ == '__main__':
    main()

考え方

参考

AtCoder ABC 163 D - Sum of Large Numbers (400 点) - けんちょんの競プロ精進記録

https://img.atcoder.jp/abc163/editorial.pdf

N+1個の0から連続する整数からk個選び、
足してできる整数の最小と最大を求める。
最大 - 最小 + 1を行うことで、各k個選択するときの
足してできる整数の数がわかる。

実装は上記の解説で理解できたが、
最大が(2 * N - k + 1) * k // 2
最小が(k - 1) * k // 2で求めることが出来る理由が理解に時間がかかった。


memo

1からnまでの和を求める公式 - 具体例で学ぶ数学

  • 最小値 が(k - 1) * k // 2になる理由
    k = 4の時、
    0 + 1 + 2 + 3 = 最小値となり、これを2倍すると、
    3 + 3 + 3 + 3 = 最小値 * 2
    3 * 4 = 最小値 * 2となる。
    3 = k - 1, 4 = kのため、
    (k - 1) * k // 2が最小値となる。

  • 最大値が(2 * N - k + 1) * k // 2になる理由
    N = 10, k = 4のとき、
    10 + 9 + 8 + 7 = 最大値となり、これを2倍すると、
    17 + 17 + 17 + 17 = 最大値 * 2
    17 * 4 = 最大値 * 2となる。
    10 + 9 + 8 + 7N + N - 1 + N - 2 + N - k + 1から、
    10 = N, 7 = N - k + 1のため、
    17 = 10 + 7 = N + N - k + 1となる。
    つまり、17 * 4 = 最大値 * 2は、
    (2N - k + 1) * k // 2 = 最大値となる。