Counting Arrays (Codeforces 1749D)
Tags: Math, combinatorics

Consider an array $$$a$$$ of length $$$n$$$ with elements numbered from $$$1$$$ to $$$n$$$. It is possible to remove the $$$i$$$-th element of $$$a$$$ if $$$gcd(a_i, i) = 1$$$, where $$$gcd$$$ denotes the greatest common divisor. After an element is removed, the elements to the right are shifted to the left by one position.An array $$$b$$$ with $$$n$$$ integers such that $$$1 \le b_i \le n - i + 1$$$ is a removal sequence for the array $$$a$$$ if it is possible to remove all elements of $$$a$$$, if you remove the $$$b_1$$$-th element, then the $$$b_2$$$-th, ..., then the $$$b_n$$$-th element. For example, let $$$a = [42, 314]$$$:  $$$[1, 1]$$$ is a removal sequence: when you remove the $$$1$$$-st element of the array, the condition $$$gcd(42, 1) = 1$$$ holds, and the array becomes $$$[314]$$$; when you remove the $$$1$$$-st element again, the condition $$$gcd(314, 1) = 1$$$ holds, and the array becomes empty.  $$$[2, 1]$$$ is not a removal sequence: when you try to remove the $$$2$$$-nd element, the condition $$$gcd(314, 2) = 1$$$ is false. An array is ambiguous if it has at least two removal sequences. For example, the array $$$[1, 2, 5]$$$ is ambiguous: it has removal sequences $$$[3, 1, 1]$$$ and $$$[1, 2, 1]$$$. The array $$$[42, 314]$$$ is not ambiguous: the only removal sequence it has is $$$[1, 1]$$$.You are given two integers $$$n$$$ and $$$m$$$. You have to calculate the number of ambiguous arrays $$$a$$$ such that the length of $$$a$$$ is from $$$1$$$ to $$$n$$$ and each $$$a_i$$$ is an integer from $$$1$$$ to $$$m$$$.

Input: The only line of the input contains two integers $$$n$$$ and $$$m$$$ ($$$2 \le n \le 3 \cdot 10^5$$$; $$$1 \le m \le 10^{12}$$$).

Output: Print one integer — the number of ambiguous arrays $$$a$$$ such that the length of $$$a$$$ is from $$$1$$$ to $$$n$$$ and each $$$a_i$$$ is an integer from $$$1$$$ to $$$m$$$. Since the answer can be very large, print it modulo $$$998244353$$$.


Sample Inputs:
2 3

4 2

4 6

1337 424242424242

Sample Outputs:
6

26

1494

119112628

Solution (Python):
(n, m) = map(int, input().split())
if m == 1:
    print(n - 1)
    exit()
M = 998244353
tot = (pow(m, n + 1, M) - 1) * pow(m - 1, -1, M) - 1
p = set([2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31])
fac = [1, 1]
for i in range(2, 37):
    fac += (fac[-1] * (i if i in p else 1),)
(i, res, temp) = (2, m, m)
while i <= min(n, 36) and fac[i] <= m:
    temp = m // fac[i] * temp % M
    res = (res + temp) % M
    i += 1
print((tot - res) % M)

Solution (Python):
(n, m) = map(int, input().split())
if m == 1:
    print(n - 1)
    exit()
M = 998244353
tot = (pow(m, n + 1, M) - 1) * pow(m - 1, -1, M) - 1
p = set([2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37])
fac = [1, 1]
for i in range(2, 38):
    fac += (fac[-1] * (i if i in p else 1),)
(i, res, temp) = (2, m, m)
while i <= min(n, 37) and fac[i] <= m:
    temp = m // fac[i] * temp % M
    res = (res + temp) % M
    i += 1
print((tot - res) % M)

Solution (Python):
from typing import List
from math import sqrt
MOD = 998244353

def sieve(n) -> List[bool]:
    is_prime = [True] * (n + 1)
    is_prime[0] = is_prime[1] = False
    for i in range(4, n + 1, 2):
        is_prime[i] = False
    for i in range(3, int(sqrt(n)) + 2, 2):
        if is_prime[i]:
            for j in range(i * i, n, 2 * i):
                is_prime[j] = False
    return is_prime

def main():
    (n, m) = map(int, input().split())
    is_prime = sieve(n)
    unambiguous = 0
    last = 1
    primes_acc = 1
    for i in range(1, n + 1):
        if primes_acc > m:
            break
        if is_prime[i]:
            primes_acc *= i
        last *= m // primes_acc
        unambiguous = (unambiguous + last) % MOD
    if m != 1:
        all = pow(m, n + 1, MOD) - 1
        all = (all * pow(m - 1, MOD - 2, MOD) - 1 + MOD) % MOD
    else:
        all = n % MOD
    ans = (all - unambiguous + MOD) % MOD
    print(ans)
if __name__ == '__main__':
    main()