HackerRank Algebra – Summing the K-N series

Challenge: Summing the K-N series

Subdomeniu: Algebra (algebra)

Scor cont: 40.0 / 40

Submission status: Accepted

Submission score: 1.0

Submission ID: 464751816

Limbaj: python3

Link challenge: https://www.hackerrank.com/challenges/summing-the-k-n-series/problem

Cerinta

You are given a sequence whose $n^{\text{th}}$ term is  
$$T_n = n^K$$
You have to evaluate the series
$$S_n = T_1 + T_2 + T_3 + \cdots + T_n$$
Find $S_n \bmod (10^9+7)$.

**Input Format**  

The first line of input contains $T$, the number of test cases.  

Each test case consists of one line containing two space-separated integers $n$ and $K$.  

**Output Format**  

For each test case, print the required answer in a line.  

**Constraints**  

$1 \le T \le 10$  
$0 \le K \le 10^3$  
$1 \le n \le 10^{16}$  

**Sample Input**

    3
    5 3
    4 2
    4 1
    
**Sample Output**  

    225
    30
    10
    
**Explanation**  

Case 1: We have $225 = 1^3 + 2^3 + 3^3 + 4^3 + 5^3$  
Case 2: We have $30  = 1^2 + 2^2 + 3^2 + 4^2$  
Case 3: We have $10  = 1^1 + 2^1 + 3^1 + 4^1$

Cod sursa

#!/bin/python3

import sys

MOD = 10 ** 9 + 7


def sum_pows(n: int, k: int) -> int:
    # f(x) = sum_{i=1..x} i^k is a polynomial of degree k+1.
    m = k + 2

    y = [0] * (m + 1)
    for i in range(1, m + 1):
        y[i] = (y[i - 1] + pow(i, k, MOD)) % MOD

    if n <= m:
        return y[n]

    fact = [1] * (m + 1)
    for i in range(1, m + 1):
        fact[i] = fact[i - 1] * i % MOD

    inv_fact = [1] * (m + 1)
    inv_fact[m] = pow(fact[m], MOD - 2, MOD)
    for i in range(m, 0, -1):
        inv_fact[i - 1] = inv_fact[i] * i % MOD

    n_mod = n % MOD

    pref = [1] * (m + 2)
    for i in range(0, m + 1):
        pref[i + 1] = pref[i] * ((n_mod - i) % MOD) % MOD

    suff = [1] * (m + 2)
    for i in range(m, -1, -1):
        suff[i] = suff[i + 1] * ((n_mod - i) % MOD) % MOD

    ans = 0
    for i in range(0, m + 1):
        num = pref[i] * suff[i + 1] % MOD
        den_inv = inv_fact[i] * inv_fact[m - i] % MOD
        if (m - i) & 1:
            den_inv = (MOD - den_inv) % MOD

        ans = (ans + y[i] * num % MOD * den_inv) % MOD

    return ans


def main():
    data = sys.stdin.buffer.read().split()
    if not data:
        return

    t = int(data[0])
    out = []
    ptr = 1
    for _ in range(t):
        n = int(data[ptr])
        k = int(data[ptr + 1])
        ptr += 2
        out.append(str(sum_pows(n, k)))

    sys.stdout.write('\n'.join(out))


if __name__ == '__main__':
    main()