Challenge: Coprime Power Sum

Subdomeniu: Number Theory (number-theory)

Scor cont: 75.0 / 75

Submission status: Accepted

Submission score: 1.0

Submission ID: 464779417

Limbaj: cpp14

Link challenge: https://www.hackerrank.com/challenges/coprime-power-sum/problem

Cerinta

Given two integers, $m$ and $k$, Alice loves to calculate their power sum using the following formula:
$$ PowerSum(m, k) \equiv \sum_{i=1}^{m} i^k $$

Bob has a set, $s$, of $n$ distinct *pairwise coprime* integers. Bob hates multiples of these integers, so he subtracts $i^k$ from Alice's power sum for each $i$ $\in$ $[1,m]$ whenever there exists at least one $j \in [1,n]$ such that $ i \bmod s_j \equiv 0$.

Alice and Bob are now confused about the final value of the power sum and decide to turn to Eve for help. Can you write a program that helps Eve solve this problem? Given $q$ queries consisting of $n$, $m$, and $k$, print the value of the power sum modulo $10^9 + 7$ on a new line for each query.

Input Format

The first line contains an integer, $q$, denoting the number of queries. The $2 \cdot q$ lines describe each query over two lines:

1. The first line contains three space-separated integers denoting the respective values of $n$ (the number of integers in Bob's set), $k$ (the exponent variable in the power sum formula), and $m$ (the upper range bound in the power sum formula). 
2. The second line contains $n$ distinct space-separated integers describing the respective elements in set $s$.

Output Format

For each query, print the resulting value of the power sum after Bob's subtraction, modulo $10^9 + 7$.

Constraints

* $1 \leq q \leq 2$  
* $1 \leq n \leq 50$  
* $0 \leq k \leq 10$  
* $1 \leq m \leq 10^{12}$  
* $1 \leq s_j \leq 10^{12}$		
* $s_i \ne s_j \text{, where } i \ne j$  	
* $gcd(s_i, s_j) \equiv 1 \text{, where } i \ne j$

Cod sursa

#include <bits/stdc++.h>
using namespace std;

static const long long MOD = 1000000007LL;

long long mod_pow(long long a, long long e) {
    a %= MOD;
    long long r = 1;
    while (e > 0) {
        if (e & 1) r = (r * a) % MOD;
        a = (a * a) % MOD;
        e >>= 1;
    }
    return r;
}

// Precompute factorials up to 105 (more than enough for k<=100 too).
struct Fact {
    vector<long long> fact, invfact;
    Fact(int n=105) {
        fact.assign(n+1, 1);
        invfact.assign(n+1, 1);
        for (int i = 1; i <= n; i++) fact[i] = fact[i-1] * i % MOD;
        invfact[n] = mod_pow(fact[n], MOD-2);
        for (int i = n; i >= 1; i--) invfact[i-1] = invfact[i] * i % MOD;
    }
} Fct;

// PowerSum via Lagrange on points 0..d where d=k+1.
struct PowerSum {
    int k;
    int d;                    // degree = k+1
    vector<long long> y;      // y[i] = sum_{t=1..i} t^k (mod MOD), i=0..d

    PowerSum(int kk=0) { init(kk); }

    void init(int kk) {
        k = kk;
        d = k + 1;
        y.assign(d + 1, 0);
        long long cur = 0;
        for (int i = 1; i <= d; i++) {
            cur = (cur + mod_pow(i, k)) % MOD;
            y[i] = cur;
        }
    }

    long long eval(long long n) const {
        if (n <= d) return y[(int)n];
        long long nm = n % MOD;

        vector<long long> pref(d+1), suf(d+1);
        // pref[i] = PI_{j=0..i} (n - j)
        pref[0] = (nm - 0 + MOD) % MOD;
        for (int i = 1; i <= d; i++)
            pref[i] = pref[i-1] * ((nm - i + MOD) % MOD) % MOD;

        // suf[i] = PI_{j=i..d} (n - j)
        suf[d] = (nm - d + MOD) % MOD;
        for (int i = d-1; i >= 0; i--)
            suf[i] = suf[i+1] * ((nm - i + MOD) % MOD) % MOD;

        long long ans = 0;
        for (int i = 0; i <= d; i++) {
            long long num = 1;
            if (i > 0) num = num * pref[i-1] % MOD;
            if (i < d) num = num * suf[i+1] % MOD;

            // denom = (i! (d-i)!)^{-1} * (-1)^{d-i}
            long long denom = Fct.invfact[i] * Fct.invfact[d-i] % MOD;
            if ((d - i) & 1) denom = (MOD - denom) % MOD;

            long long term = y[i] * num % MOD * denom % MOD;
            ans += term;
            if (ans >= MOD) ans -= MOD;
        }
        return ans;
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int Q;
    cin >> Q;
    while (Q--) {
        int n, k;
        long long m;
        cin >> n >> k >> m;

        vector<long long> s;
        s.reserve(n);
        bool has_one = false;

        for (int i = 0; i < n; i++) {
            long long x;
            cin >> x;
            if (x == 1) has_one = true;
            if (x <= m) s.push_back(x); // ignore x>m (no multiples <=m)
        }

        if (has_one) {
            cout << 0 << "\n";
            continue;
        }
        if (s.empty()) {
            // No forbidden numbers effectively
            PowerSum ps(k);
            cout << ps.eval(m) << "\n";
            continue;
        }

        sort(s.begin(), s.end());
        int N = (int)s.size();

        PowerSum ps(k);

        vector<long long> pk(N);
        for (int i = 0; i < N; i++)
            pk[i] = mod_pow(s[i], k);

        // memo[i] : unordered_map<x, value> for F(i,x) where i is 0..N
        vector<unordered_map<long long,long long>> memo(N+1);

        function<long long(int,long long)> dfs = [&](int i, long long x) -> long long {
            if (x == 0) return 0;
            if (i == 0) return ps.eval(x);

            auto &mp = memo[i];
            auto it = mp.find(x);
            if (it != mp.end()) return it->second;

            long long ans;
            long long si = s[i-1];

            if (si > x) {
                ans = dfs(i-1, x);
            } else {
                long long A = dfs(i-1, x);
                long long B = dfs(i-1, x / si);
                ans = (A - pk[i-1] * B) % MOD;
                if (ans < 0) ans += MOD;
            }
            mp[x] = ans;
            return ans;
        };

        cout << dfs(N, m) << "\n";
    }
    return 0;
}
HackerRank Number Theory – Coprime Power Sum