3123 - summy

Cerința

Se dau n şi k numere naturale. Calculați suma ∑ni=1ik.

Date de intrare

Se dau n şi k numere naturale. Calculați suma ∑ni=1ik.

Date de ieșire

Programul va afișa pe ecran valoarea sumei ∑ni=1ik, modulo 1.000.000.007.

Restricții și precizări

• 1 ≤ n ≤ 100.000 şi 1 ≤ k ≤ 1.000.000.000 pentru 70% din teste
• 1 ≤ n ≤ 1.000.000.000 şi 1 ≤ k ≤ 100.000 pentru 30% din teste

Exemplul 1:

Intrare

5 3

Ieșire

225

Exemplul 2:

Intrare

5 3

Ieșire

7 8
7907396

Rezolvare

MOD = 1000000007

def gcd(a, b):
    if b == 0:
        return (1, 0)
    else:
        x, y = gcd(b, a % b)
        return (y, x - y * (a // b))

def invers(x):
    inv, _ = gcd(x, MOD)
    if inv <= 0:
        inv = MOD + inv % MOD
    return inv

def pu(b, e):
    if e == 0:
        return 1
    elif e % 2 == 0:
        x = pu(b, e // 2)
        return (x * x) % MOD
    else:
        return (b * pu(b, e - 1)) % MOD

def calculate_for_small_n(n, k):
    sol = 0
    for i in range(1, n + 1):
        sol = (sol + pu(i, k)) % MOD
    return sol

def verifica_restrictii(n, k):
    return 1 <= n <= 100000 and 1 <= k <= 1000000000

def main():
    n, k = map(int, input().split())
    if verifica_restrictii(n, k):
        if n <= 100000:
            print(calculate_for_small_n(n, k))
        else:
            suma = 1
            numarator = 1
            numitor = 1
            for i in range(2, k + 3):
                numarator = (numarator * (n - i)) % MOD
                numitor = (numitor * (1 - i)) % MOD
            if numitor < 0:
                numitor += MOD
            sol = (numarator * invers(numitor)) % MOD
            for i in range(2, k + 3):
                suma = (suma + pu(i, k)) % MOD
                numarator = (numarator * (n - i + 1)) % MOD
                numarator = (numarator * invers(n - i)) % MOD
                numitor = (numitor * (i - 1)) % MOD
                numitor = (-numitor * invers(k + 3 - i)) % MOD
                if numitor < 0:
                    numitor += MOD
                y = (numarator * invers(numitor)) % MOD
                sol = (sol + suma * y) % MOD
            print(sol)
    else:
        print("Datele nu corespund restrictiilor impuse")

if __name__ == "__main__":
    main()

Explicatie

∑5i=1i*3=1**3+2**3+3**3+4**3+5**3=225