haar_lib/math/
sum_of_exponential_times_polynomial_limit.rs

1//! $\sum_{i=0}^{\infty} r^ii^d$
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/sum_of_exponential_times_polynomial_limit>
5
6use crate::num::ff::*;
7
8/// $\sum_{i=0}^{\infty} r^ii^d$
9///
10/// **Time Complexity** $O(d \log p)$
11pub fn sum_of_exponential_times_polynomial_limit<Modulo: FF>(
12    r: Modulo::Element,
13    d: u64,
14    m: Modulo,
15) -> Modulo::Element
16where
17    Modulo::Element: FFElem + Copy,
18{
19    let mut ret = m.from_u64(0);
20    let mut r_pow = m.from_u64(1);
21    let mut s = vec![m.from_u64(0); d as usize + 1];
22    let mut invs = vec![m.from_u64(0); d as usize + 2];
23    let p = m.modulo();
24
25    invs[1] = m.from_u64(1);
26
27    for i in 2..=d + 1 {
28        invs[i as usize] = m.from_u64(p as u64 / i) * -invs[(p as u64 % i) as usize];
29    }
30
31    for i in 0..=d as usize {
32        if i > 0 {
33            let x = s[i - 1];
34            s[i] += x;
35        }
36        s[i] += m.from_u64(i as u64).pow(d) * r_pow;
37        r_pow *= r;
38    }
39
40    let mut t = m.from_u64(1);
41
42    for i in 0..=d {
43        ret += t * s[(d - i) as usize];
44        t *= invs[i as usize + 1] * -r * m.from_u64(d + 1 - i);
45    }
46
47    ret * (m.from_u64(1) - r).pow(d + 1).inv()
48}