The statement:
Given three integers $$$n, k, p$$$, $$$(1 \leq k \leq n < p)$$$.
Count the number of array $$$a[]$$$ of size $$$k$$$ that satisfied
- $$$1 \leq a_1 < a_2 < \dots < a_k \leq n$$$
- $$$a_i \times a_j$$$ is perfect square $$$\forall 1 \leq i < j \leq n$$$
Since the number can be big, output it under modulo $$$p$$$.
For convenient, you can assume $$$p$$$ is a large constant prime $$$10^9 + 7$$$
Yet you can submit the problem for $$$k = 3$$$ here.
Extra Tasks
Solved A: Can we also use phi function or something similar to solve for $$$k = 3$$$ in $$$O(\sqrt{n})$$$ ?
Solved B: Can we also use phi function or something similar to solve for general $$$k$$$ in $$$O(\sqrt{n})$$$ ?
Solved C: Can we also solve the problem where there can be duplicate: $$$a_i \leq a_j\ (\forall\ i < j)$$$ and no longer $$$a_i < a_j (\forall\ i < j)$$$ ?
Solved D: Can we solve the problem where there is no restriction between $$$k, n, p$$$ ?
E: Can we solve for negative integers, whereas $$$-n \leq a_1 < a_2 < \dots < a_k \leq n$$$ ?
F: Can we solve for a specific range, whereas $$$L \leq a_1 < a_2 < \dots < a_k \leq R$$$ ?
G: Can we solve for cube product $$$a_i \times a_j \times a_k$$$ effectively ?
H: Can we solve if it is given $$$n$$$ and queries for $$$k$$$ ?
I: Can we solve if it is given $$$k$$$ and queries for $$$n$$$ ?
J: Can we also solve the problem where there are no order: Just simply $$$1 \leq a_i \leq n$$$ ?
K: Can we solve for $$$q$$$-product $$$a_{i_1} \times a_{i_2} \times \dots \times a_{i_q} = x^q$$$ (for given constant $$$q$$$) ?
M: Given $$$0 \leq \delta \leq n$$$, can we also solve the problem when $$$1 \leq a_1 \leq a_1 + \delta + \leq a_2 \leq a_2 + \delta \leq \dots \leq a_k \leq n$$$ ?
*Marked as solved only if tested with atleast $$$10^6$$$ queries
Solution for k = 1
The answer just simply be $$$n$$$
Solution for k = 2
Algorithm
We need to count the number of pair $$$(a, b)$$$ that $$$1 \leq a < b \leq n$$$ and $$$a \times b$$$ is perfect square.
Every positive integer $$$x$$$ can be represent uniquely as $$$x = u \times p^2$$$ for some positive integer $$$u, p$$$ and $$$u$$$ as small as possible ($$$u$$$ is squarefree number).
Let represent $$$x = u \times p^2$$$ and $$$y = v \times q^2$$$ (still, minimum $$$u$$$, $$$v$$$ ofcourse).
We can easily proove that $$$x \times y$$$ is a perfect square if and if only $$$u = v$$$.
So for a fixed squarefree number $$$u$$$. You just need to count the number of ways to choose $$$p^2$$$.
The answer will be the sum of such ways for each fixed $$$u$$$.
Implementation
Implementation using factorizationvector<int> prime; /// prime list
vector<int> lpf; /// Lowest prime factor, lpf[x] is smallest prime divisor of x
void sieve(int lim = LIM) /// O(n)
{
prime.assign(1, 2);
lpf.assign(lim + 1, 2);
lpf[1] = 1; /// For easier calculation but can cause inf loops
for (int i = 3; i <= lim; i += 2) {
if (lpf[i] == 2) prime.push_back(lpf[i] = i);
for (int j = 0; j < sz(prime) && prime[j] <= lpf[i] && prime[j] * i <= lim; ++j)
lpf[prime[j] * i] = prime[j];
}
}
/// mask(x) is smallest positive number that mask(x) * x is a perfect square
int getMask(int x) /// O(log n)
{
int mask = 1;
while (x > 1) {
int p = lpf[x], f = 0;
do x /= p, f++; while (p == lpf[x]);
if (f & 1) mask *= p; /// if current power is odd, we mutiple mask with current prime
}
return mask;
}
int cnt[LIM];
int magic(int n) /// O(n log max(a))
{
memset(cnt, 0, sizeof(cnt[0]) * (n + 1));
ll res = 0;
for (int a = 1; a <= n; ++a) /// Check all cases of a
res += cnt[getMask(a)]++;
res %= MOD;
return res;
}
int main() /// O(n log max(a))
{
int n;
cin >> n;
sieve(n + 500);
cout << magic(n);
return 0;
}
Implementation 1int solve(int n)
{
memset(is_squarefree, true, sizeof(is_squarefree[0]) * (n + 1));
long long res = 0;
for (int i = 1, j; i <= n; ++i) if (is_squarefree[i])
{
for (j = 1; i * j * j <= n; ++j)
is_squarefree[i * j * j] = false;
res += 1LL * (j - 1) * (j - 2) / 2;
}
res %= MOD;
return res;
}
Implementation 2int solve(int n)
{
memset(is_squarefree, true, sizeof(is_squarefree[0]) * (n + 1));
long long res = 0;
for (int i = 1; i <= n; ++i) if (is_squarefree[i])
{
for (int j = 1; i * j * j <= n; ++j)
{
is_squarefree[i * j * j] = false;
res += j - 1;
}
}
res %= MOD;
return res;
}
Implementation related to Möbius function/// Linear Sieve
vector<bool> isPrime; /// Characteristic function = A010051
vector<int> prime; /// Prime list = A000040
vector<int> lpf; /// Lowest prime factor = A020639
vector<int> mu; /// Mobius = A008683
vector<int> phi; /// Euler totient phi = A000010
void sieve(int n)
{
if (n < 1) return ;
/// Extension part
mu.assign(n + 1, 1);
phi.assign(n + 1, 1);
/// Main part
prime.clear();
lpf.assign(n + 1, 0);
isPrime.assign(n + 1, true);
isPrime[0] = isPrime[1] = false;
for (int x = 2; x <= n; ++x)
{
if (isPrime[x]) /// Func[Prime]
{
lpf[x] = x;
mu[x] = -1;
phi[x] = x - 1;
prime.push_back(x);
}
for (int p : prime) /// Func[Prime * X] <- Func[Prime]
{
if (p > lpf[x] || x * p > n) break;
isPrime[x * p] = 0;
lpf[x * p] = p;
if (lpf[x] == p)
{
mu[x * p] = 0;
phi[x * p] = phi[x] * p;
}
else
{
mu[x * p] = -mu[x];
phi[x * p] = phi[x] * phi[p];
}
}
}
}
long long solve(int n)
{
sieve(n + 1);
long long res = 0;
for (int x = 1; x <= n; ++x) if (mu[x])
{
int t = sqrt(n / x);
res += 1LL * t * (t - 1) / 2;
}
res %= MOD;
return res;
}
Complexity
So about the complexity....
For the implementation using factorization, it is $$$O(n \log n)$$$.
Hint 2Precalculating smallest prime factor for each $$$n$$$.
ProofSince you known the smallest prime factor for each $$$n$$$, you can factorize any number in $$$O(\log n)$$$.
There are $$$n$$$ number total. Hence the calculating part take $$$O(n \log n)$$$.
The sieve only cost you $$$O(n)$$$ give the total complexity of $$$O(n \log n)$$$.
BonusThis just a clickbait. Have a nice day.
Real Bonus The true BonusYou can optimize it to solve in $$$O(n \log \log n)$$$ though it is not that necessary
For the 2 implementations below, the complexity is linear.
Hint 1Each number will be calculated once.
Hint 2Though some numbers might be visited many times. It is actually armotized to constant.
Hint 3Let remind you about something called the sum of inversion of squares.
Hint 4Just ignore those squarefree. A normal loop through it is still linear.
Proof$$$O\left(\underset{\text{squarefree } p \leq n}{\Large \Sigma} \left \lfloor {\Large \frac{n}{p^2}} \right \rfloor \right) = O\left(\underset{p \leq n}{\Large \Sigma} {\Large \frac{n}{p^2}} \right) = O\left(n \times {\Large \Sigma} {\Large \frac{1}{p^2}} \right) = O\left(n \times {\Large \frac{\pi^2}{6}} \right) = O(n)$$$.
For the last implementation, the complexity is Linear
Hint 1It depends on the sieve as the calculation is already linear.
Hint 2Each number can be represent as $$$x = lpf[x] \times k$$$ where $$$lpf[x]$$$ is lowest prime factor of $$$x$$$ and $$$k$$$ is an integer
ProofEach $$$x$$$ can be represented uniquely as the form $$$x = lpf[x] \times k$$$.
So you can make a sieve of linear, where each number is visisted once or atmost $$$O(1)$$$
Solution for general k
Using the same logic above, we can easily solve the problem.
Now you face up with familliar binomial coefficient problem
This implementation here is using the assumption of $$$p$$$ prime and $$$p > max(n, k)$$$
You can still solve the problem for squarefree $$$p$$$ using lucas and CRT
Yet just let things simple as we only focus on the counting problem, we will assume $$$p$$$ is a large constant prime.
O(n) solutionconst int LIM = 1e7 + 17;
const int MOD = 1e9 + 7;
int fact[LIM + 1]; /// factorial: fact[n] = n!
int invs[LIM + 1]; /// inverse modular: invs[n] = n^(-1)
int tcaf[LIM + 1]; /// inverse factorial: tcaf[n] = (n!)^(-1)
void precal_nck(int n = LIM)
{
fact[0] = fact[1] = 1;
invs[0] = invs[1] = 1;
tcaf[0] = tcaf[1] = 1;
for (int i = 2; i <= n; ++i)
{
fact[i] = (1LL * fact[i - 1] * i) % MOD;
invs[i] = MOD - 1LL * (MOD / i) * invs[MOD % i] % MOD;
tcaf[i] = (1LL * tcaf[i - 1] * invs[i]) % MOD;
}
}
int nCk(int n, int k)
{
k = min(k, n - k);
if (k < 0) return 0;
long long res = fact[n];
res *= tcaf[k]; res %= MOD;
res *= tcaf[n - k]; res %= MOD;
return res;
}
int solve(int n, int k)
{
memset(is_squarefree, true, sizeof(is_squarefree[0]) * (n + 1));
precal_nck(n);
long long res = 0;
for (int i = 1, j; i <= n; ++i) if (is_squarefree[i])
{
for (j = 1; i * j * j <= n; ++j)
is_squarefree[i * j * j] = false;
res += nCk(j - 1, k);
}
res %= MOD;
return res;
}
A better solution for k = 2
Idea
In the above approach, we fix $$$u$$$ as a squarefree and count $$$p^2$$$.
But what if I fix $$$p^2$$$ to count $$$u$$$ instead ?
Yet you can see that the first loop now is $$$O(\sqrt{n})$$$, but it will still $$$O(n)$$$ total because of the second loop
Swap for loop implementationint solve(int n)
{
memset(is_squarefree, true, sizeof(is_squarefree[0]) * (n + 1));
long long res = 0;
int t = sqrt(n);
while (t * t < n) ++t;
while (t * t > n) --t;
for (int j = t; j > 1; --j)
{
for (int i = 1; i * j * j <= n; ++i)
{
if (!used[i * j * j])
{
used[i * j * j] = true;
res += j - 1;
}
}
}
res %= MOD;
return res;
}
Approach
Let $$$f(n)$$$ is the number of pair $$$(a, b)$$$ that $$$1 \leq a < b \leq n$$$ and $$$(a, b, n)$$$ is a three-term geometric progression.
Let $$$g(n)$$$ is the number of pair $$$(a, b)$$$ that $$$1 \leq a \leq b \leq n$$$ and $$$(a, b, n)$$$ is a three-term geometric progression.
Let $$$F(n) = \overset{n}{\underset{p=1}{\Large \Sigma}} f(p)$$$.
But why do we need these functions anywayWell, you see, since $$$(a, b, n)$$$ form a three-term geometric progression then you have $$$b^2 = an$$$.
It is not hard to proove that $$$F(N)$$$ will be our answer as we count over all possible squarefree $$$u$$$ for every fixed $$$p^2$$$.
We use $$$g(n)$$$ because its beautiful property, which make the calculation a lot easier.
So it is no hard to prove that $$$g(n) = f(n) + 1$$$.
This interesting sequence $$$g(n)$$$ is A000188, having many properties, such as
- Number of solutions to $$$x^2 \equiv 0 \pmod n$$$.
- Square root of largest square dividing $$$n$$$.
- Max $$$gcd \left(d, \frac{n}{d}\right)$$$ for all divisor $$$d$$$.
Well, to make the problem whole easier, I gonna skip all the proofs to use this property (still, you can use the link in the sequence for references).
$$$g(n) = \underset{d^2 | n}{\Large \Sigma} \phi(d)$$$.
From this property, we can solve the problem in $$$O(\sqrt{n})$$$.
Hint 1$$$F(N) = \overset{n}{\underset{p=1}{\Large \Sigma}} f(p)$$$.
and
$$$f(n) = g(n) + 1$$$.
Hint 2We have
$$$F(n)$$$
$$$= \overset{n}{\underset{p=2}{\Large \Sigma}} g(p) - n$$$
$$$= \overset{n}{\underset{p=2}{\Large \Sigma}} \underset{d^2 | p}{\Large \Sigma} \phi(d) - n$$$
Hint 3Break the loop, make it simpler.
Hint 4Reduce a for loop to constant.
Solution$$$F(N)$$$
$$$= \overset{n}{\underset{p=1}{\Large \Sigma}} f(p)$$$
$$$= -n + \overset{n}{\underset{p=2}{\Large \Sigma}} g(p)$$$
$$$= -n + \overset{n}{\underset{p=2}{\Large \Sigma}} \underset{d^2 | p}{\Large \Sigma} \phi(d)$$$
$$$= -n + \overset{\left \lfloor \sqrt{n} \right \rfloor}{\underset{p=1}{\Large \Sigma}} \underset{1 \leq u \times p^2 \leq n}{\Large \Sigma} \phi(p)$$$
$$$= -n + \overset{\left \lfloor \sqrt{n} \right \rfloor}{\underset{p=1}{\Large \Sigma}} \phi(p) \times \left \lfloor \frac{n}{p^2} \right \rfloor$$$
$$$= \overset{\left \lfloor \sqrt{n} \right \rfloor}{\underset{p=2}{\Large \Sigma}} \phi(p) \times \left \lfloor \frac{n}{p^2} \right \rfloor$$$
Well, you can sieve $$$\phi(p)$$$ for all $$$2 \leq p \leq \sqrt{n}$$$ in $$$O\left(\sqrt{n} \log \log \sqrt{n} \right)$$$ or improve it with linear sieve to $$$O(\sqrt{n})$$$.
Therefore you can solve the whole problem in $$$O(\sqrt{n})$$$.
Yet this paper also takes you to something similar.
Implementation
O(sqrt n log log sqrt n) solution#include <iostream>
#include <cstring>
#include <numeric>
#include <cmath>
using namespace std;
const int MOD = 1e9 + 7;
const int LIM = 1e7 + 17;
const int SQRT_LIM = ceil(sqrt(LIM) + 1) + 1;
int euler[SQRT_LIM];
void sieve_phi(int n)
{
iota(euler, euler + n + 1, 0);
for (int x = 2; x <= n; x++) if (euler[x] == x)
for (int j = x; j <= n; j += x)
euler[j] -= euler[j] / x;
}
int solve(int n)
{
sieve_phi(ceil(sqrt(n) + 1) + 1);
long long res = 0;
for (int p = 2; p * p <= n; ++p)
res += 1LL * euler[p] * (n / (p * p));
res %= MOD;
return res;
}
int main()
{
ios::sync_with_stdio(false);
cin.tie(NULL);
int n;
cin >> n;
cout << solve(n);
return 0;
}
O(sqrt) solution#include <iostream>
#include <cstring>
#include <numeric>
#include <vector>
#include <cmath>
using namespace std;
const int MOD = 1e9 + 7;
vector<int> lpf;
vector<int> prime;
vector<int> euler;
void linear_sieve_phi(int n)
{
lpf.assign(n + 1, 0);
euler.assign(n + 1, 1);
for (int x = 2; x <= n; ++x)
{
if (lpf[x] == 0)
{
prime.push_back(lpf[x] = x);
euler[x] = x - 1;
}
for (int i = 0; i < prime.size() && x * prime[i] <= n; ++i)
{
lpf[x * prime[i]] = prime[i];
if (x % prime[i] == 0) {
euler[x * prime[i]] = euler[x] * prime[i];
break;
}
euler[x * prime[i]] = euler[x] * euler[prime[i]];
}
}
}
int solve(int n)
{
linear_sieve_phi(ceil(sqrt(n) + 1) + 1);
long long res = 0;
for (int p = 2; p * p <= n; ++p)
res += 1LL * euler[p] * (n / (p * p));
res %= MOD;
return res;
}
int main()
{
ios::sync_with_stdio(false);
cin.tie(NULL);
int n;
cin >> n;
cout << solve(n);
return 0;
}
A better solution for general k
Extra task A, B
Algorithm
As what clyring decribed here
Let $$$f_k(n)$$$ is the number of set $$$(a_1, a_2, \dots, a_k, n)$$$ that $$$1 \leq a_1 < a_2 < \dots < a_k \leq n$$$ and $$$(a_1, a_2, \dots, a_k, n)$$$ is a $$$(k+1)$$$-term geometric progression.
Let $$$g_k(n)$$$ is the number of set $$$(a_1, a_2, \dots, a_k, n)$$$ that $$$1 \leq a_1 \leq a_2 \leq \dots \leq a_k \leq n$$$ and $$$(a_1, a_2, \dots, a_k, n)$$$ is a $$$(k+1)$$$-term geometric progression.
Let $$$F_k(n) = \overset{n}{\underset{p=1}{\Large \Sigma}} f_k(p)$$$.
Let $$$s_k(n)$$$ is the number of way to choose $$$p^2$$$ among those $$$k$$$ numbers when you fix squarefree $$$u$$$ (though we are doing in reverse).
The formula$$$s_k(n) = C^{n-1}_{k-1}$$$
$$$f_k(n) = F_k(n) - F_k(n-1) = s_k(g_k(n)) = \underset{d^2 | n}{\Large \Sigma} (\mu \times s_k)(d) = \underset{d^2 | n}{\Large \Sigma} \underset{p | d} \mu(p) \times s_k\left(\frac{d}{p}\right)$$$
$$$F_k(n) = \overset{n}{\underset{p=1}{\Large \Sigma}} f_k(p)$$$
$$$= F_k(0) + \overset{\left \lfloor \sqrt{n} \right \rfloor}{\underset{d = 1}{\Large \Sigma}}(\mu \times s_k)(d) \times \left \lfloor \frac{n}{d^2} \right \rfloor$$$
$$$= \overset{\left \lfloor \sqrt{n} \right \rfloor}{\underset{d = 1}{\Large \Sigma}} \left \lfloor \frac{n}{d^2} \right \rfloor \times \left ( \underset{p | d}{\Large \Sigma} \mu(p) \times s_k\left(\frac{d}{p}\right) \right)$$$
Implementation
O(sqrt n log sqrt n)#include <iostream>
#include <cstring>
#include <numeric>
#include <vector>
#include <cmath>
using namespace std;
const int LIM = 5e6 + 56;
const int SQRT_LIM = ceil(sqrt(LIM) + 1) + 1;
const int MOD = 1e9 + 7;
/// Precalculating factorials under prime modulo
int fact[SQRT_LIM + 10]; /// fact[n] = n!
int invs[SQRT_LIM + 10]; /// invs[n] = n^(-1)
int tcaf[SQRT_LIM + 10]; /// tcaf[n] = (n!)^(-1)
void precal_nck(int n = SQRT_LIM)
{
fact[0] = fact[1] = 1;
invs[0] = invs[1] = 1;
tcaf[0] = tcaf[1] = 1;
for (int i = 2; i <= n; ++i)
{
fact[i] = (1LL * fact[i - 1] * i) % MOD;
invs[i] = MOD - 1LL * (MOD / i) * invs[MOD % i] % MOD;
tcaf[i] = (1LL * tcaf[i - 1] * invs[i]) % MOD;
}
}
/// Calculating binomial coefficient queries
int nck(int n, int k)
{
k = min(k, n - k);
if (k < 0) return 0;
long long res = fact[n];
res *= tcaf[k]; res %= MOD;
res *= tcaf[n - k]; res %= MOD;
return res;
}
/// Linear Sieve
vector<int> prime; /// prime list = A000040
bool isPrime[SQRT_LIM + 10]; /// characteristic function = A010051
int lpf[SQRT_LIM + 10]; /// lowest prime factor = A020639
int mu[SQRT_LIM + 10]; /// mobius = A008683
void linear_sieve(int n)
{
if (n < 1) return ;
/// Extension Sieve || You can add something more
memset(lpf, 0, sizeof(lpf[0]) * (n + 1));
fill_n(mu, n + 1, 1);
/// Main Sieve || Without this, you barely able to achive linear complexity
prime.clear();
prime.reserve(n / log(n - 1));
memset(isPrime, true, sizeof(isPrime[0]) * (n + 1));
isPrime[0] = isPrime[1] = false;
for (int x = 2; x <= n; ++x) /// For each number
{
if (isPrime[x]) /// Func[Prime]
{
mu[x] = -1;
lpf[x] = x;
prime.push_back(x);
}
for (int p : prime) /// Func[Prime * X] <- Func[Prime]
{
if (p > lpf[x] || x * p > n) break;
isPrime[x * p] = 0;
lpf[x * p] = p;
mu[x * p] = (lpf[x] == p) ? 0 : -mu[x];
}
}
}
/// Divisor sieve
vector<int> divisors[SQRT_LIM];
void precal_div(int n) /// O(n log n)
{
for (int u = n; u >= 1; --u)
{
divisors[u].clear();
for (int v = u; v <= n; v += u)
divisors[v].push_back(u);
}
}
/// Solving for n, k
long long solve(int n, int k)
{
/// We only care for d that 1 <= d <= sqrt(n)
int t = ceil(sqrt(n) + 1) + 1;
linear_sieve(t);
precal_nck(t);
precal_div(t);
long long res = 0;
for (int d = 1; d * d <= n; ++d) /// For each fixed p^2
{
long long sum = 0;
for (int p : divisors[d]) /// For each (p | d)
sum += mu[d / p] * nck(p - 1, k - 1);
sum %= MOD;
res += sum * (n / (d * d));
}
res %= MOD;
return res;
}
int main()
{
ios::sync_with_stdio(false);
cin.tie(NULL);
/// Assumming constant p = 10^9 + 7
int n, k;
cin >> n >> k;
cout << solve(n, k);
return 0;
}
O(sqrt log log sqrt n)vector<int> prime; /// prime list = A000040
bool isPrime[SQRT_LIM + 10]; /// characteristic function = A010051
int lpf[SQRT_LIM + 10]; /// lowest prime factor = A020639
void linear_sieve(int n)
{
if (n < 1) return ;
prime.clear();
prime.reserve(n / log(n - 1));
memset(lpf, 0, sizeof(lpf[0]) * (n + 1));
memset(isPrime, true, sizeof(isPrime[0]) * (n + 1));
isPrime[0] = isPrime[1] = false;
for (int x = 2; x <= n; ++x)
{
if (isPrime[x]) /// Func[Prime]
{
lpf[x] = x;
prime.push_back(x);
}
for (int p : prime) /// Func[Prime * X] <- Func[Prime]
{
if (p > lpf[x] || x * p > n) break;
isPrime[x * p] = 0;
lpf[x * p] = p;
}
}
}
long long res[SQRT_LIM + 10];
long long solve(int n, int k)
{
int t = ceil(sqrt(n) + 1) + 1;
linear_sieve(t);
precal_nck(t);
memset(res, 0, sizeof(res[0]) * (t + 1));
for (int d = 1; d * d <= n; ++d)
res[d] = nck(d - 1, k - 1);
for (int p : prime)
for (int d = t / p; d > 0; --d)
res[d * p] -= res[d];
long long ans = 0;
for (int d = 1; d * d <= n; ++d)
ans += res[d] * (n / (d * d));
ans %= MOD;
return ans;
}
Complexity
The complexity of the first implementation is $$$O(\sqrt{n} \log \sqrt{n})$$$
Hint 1The complexity depends on Möbius function, traverse over all divisor of each number and calculating binomial coefficients
Hint 2We only care about $$$d$$$ that $$$1 \leq d \leq \sqrt{n}$$$
Hint 3We can use factorial formula for binomial coefficients as $$$1 \leq k \leq n < p$$$
ProofAs we only use $$$d$$$ that $$$d^2 \leq n$$$, let $$$t = \sqrt{n}$$$ for convenient
We can use linear sieve for $$$\mu{d}$$$ in $$$O(t)$$$
We can use divisor sieve for finding all divisors $$$p$$$ for each $$$v$$$ in $$$O(t \log t)$$$
We can calculate the result in $$$O(t \log t)$$$
Therefore the overall complexity is $$$O(t \log t)$$$
The complexity of the second implementation is $$$O(\sqrt{n} \log \log \sqrt{n})$$$
Hint 1The complexity depends on sieve, calculating result and calculating binomial coefficients
ProofThe sieve and binomial coefficients is obviously $$$O(t)$$$ complexity for $$$t = \sqrt{n}$$$
The calculating part is just simply related to the sum of inversion of primes
Therefore we achived $$$O(t \log t)$$$ complexity for $$$t = \sqrt{n}$$$
Solution for duplicates elements in array
Extra task C
Idea
It is no hard to proove that we can use the same algorithm as described in task A, B or in original task.
HintLiterally what we do up to now just optimize the counting part. The counting, which is the core is as the same as original.
ProofThe counting part is only about the number of valid sequences bounded by $$$n$$$, or you can say, the number of possible $$$p^2$$$ you can choose when you fix squarefree number $$$u$$$.
And those stuffs above $$$\mu(n)$$$, $$$\phi(n)$$$, sieving or anything else are just optimizations.
Using the same algorithm, the core of calculating is to find out the number of non-decreasing integer sequence of size $$$k$$$ where numbers are in $$$[1, n]$$$.
The formula is$$$\binom{n + k - 1}{k}$$$
and when you fix one number as a similar idea:
$$$s_k(n) = \binom{n + k - 2}{k - 1}$$$
Can you proove it ?
Hint 1It is standard stars and bars problem.
Hint 2Instead of thinking the number of way to select $$$a_i$$$.
You can think of the number of value $$$a_i$$$ you have used.
Hint 3Let $$$x_v$$$ is the number of value $$$v$$$ you used, or the number of $$$a_i = v$$$ in this sequence.
Each sequence can be represented uniquely as a sequence of $$$x_1, x_2, \dots, x_n$$$.
Since we selected $$$k$$$ in total we have $$$x_1 + x_2 + \dots + x_n = k$$$.
ProofLet $$$x_v$$$ is the number of value $$$v$$$ you used, or the number of $$$a_i = v$$$ in this sequence.
For each box $$$v = 1 \dots n$$$, draw $$$x_v$$$ stars ⭐.
You can make $$$k - 1$$$ cut between these box (not cut into the box), and split into $$$k$$$ intervals.
The sum of each interval is the value of selected element.
Since there are $$$n + k - 1$$$ boxes, and you use $$$k - 1$$$ cuts, there are the total of $$$\binom{n + k - 1}{k}$$$ different ways.
Now it is done, just that it
The idea is the same as what clyring described here but represented in the other way
Implementation
O(n^3 log n) solutionlong long brute(int n)
{
set<int> S;
for (int i = 1; i <= n; ++i)
S.insert(i * i);
long long res = 0;
for (int i = 1; i <= n; ++i)
for (int j = i; j <= n; ++j)
if (S.count(i * j))
for (int k = j; k <= n; ++k)
if (S.count(i * k))
if (S.count(j * k))
++res;
return res;
}
O(n) solution
int fact[SQRT_LIM + 10];
int invs[SQRT_LIM + 10];
int tcaf[SQRT_LIM + 10];
void precal_nck(int n = SQRT_LIM)
{
fact[0] = fact[1] = 1;
invs[0] = invs[1] = 1;
tcaf[0] = tcaf[1] = 1;
for (int i = 2; i <= n; ++i)
{
fact[i] = (1LL * fact[i - 1] * i) % MOD;
invs[i] = MOD - 1LL * (MOD / i) * invs[MOD % i] % MOD;
tcaf[i] = (1LL * tcaf[i - 1] * invs[i]) % MOD;
}
}
int nCk(int n, int k)
{
k = min(k, n - k);
if (k < 0) return 0;
long long res = fact[n];
res *= tcaf[k]; res %= MOD;
res *= tcaf[n - k]; res %= MOD;
return res;
}
bool is_squarefree[LIM];
int solve(int n, int k)
{
memset(is_squarefree, true, sizeof(is_squarefree[0]) * (n + 1));
precal_nck(2 * n + 1);
long long res = 0;
for (int i = 1, j; i <= n; ++i) if (is_squarefree[i])
{
for (j = 1; i * j * j <= n; ++j)
is_squarefree[i * j * j] = false;
res += nCk(k + j - 2, k);
}
res %= MOD;
return res;
}
O(sqrt n log log sqrt n) solutionint fact[SQRT_LIM + 10];
int invs[SQRT_LIM + 10];
int tcaf[SQRT_LIM + 10];
void precal_nck(int n = SQRT_LIM)
{
fact[0] = fact[1] = 1;
invs[0] = invs[1] = 1;
tcaf[0] = tcaf[1] = 1;
for (int i = 2; i <= n; ++i)
{
fact[i] = (1LL * fact[i - 1] * i) % MOD;
invs[i] = MOD - 1LL * (MOD / i) * invs[MOD % i] % MOD;
tcaf[i] = (1LL * tcaf[i - 1] * invs[i]) % MOD;
}
}
int nck(int n, int k)
{
k = min(k, n - k);
if (k < 0) return 0;
long long res = fact[n];
res *= tcaf[k]; res %= MOD;
res *= tcaf[n - k]; res %= MOD;
return res;
}
vector<int> prime; /// prime list = A000040
bool isPrime[SQRT_LIM + 10]; /// characteristic function = A010051
int lpf[SQRT_LIM + 10]; /// lowest prime factor = A020639
void linear_sieve(int n)
{
if (n < 1) return ;
prime.clear();
prime.reserve(n / log(n - 1));
memset(lpf, 0, sizeof(lpf[0]) * (n + 1));
memset(isPrime, true, sizeof(isPrime[0]) * (n + 1));
isPrime[0] = isPrime[1] = false;
for (int x = 2; x <= n; ++x)
{
if (isPrime[x]) /// Func[Prime]
{
lpf[x] = x;
prime.push_back(x);
}
for (int p : prime) /// Func[Prime * X] <- Func[Prime]
{
if (p > lpf[x] || x * p > n) break;
isPrime[x * p] = 0;
lpf[x * p] = p;
}
}
}
long long res[SQRT_LIM + 10];
long long solve(int n, int k)
{
int t = ceil(sqrt(n) + 1) + 1;
linear_sieve(t);
precal_nck(t);
memset(res, 0, sizeof(res[0]) * (t + 1));
for (int d = 1; d * d <= n; ++d)
res[d] = nck(d + k - 2, k - 1);
for (int p : prime)
for (int d = t / p; d > 0; --d)
res[d * p] -= res[d];
long long ans = 0;
for (int d = 1; d * d <= n; ++d)
ans += res[d] * (n / (d * d));
ans %= MOD;
return ans;
}
Complexity
Well, if you are here then I bet you a discord nitro that you dont need more proofs, lol.
Solution when there are no restriction between k, n, p
Extra task D
Idea
So first of all, the result do depend on how you calculate binomial coefficient but they are calculated independently even if you can somehow manage to use the for loop of binomial coeffient go first.
Therefore even if there is no restriction between $$$k, n, p$$$, the counting part and the algorithm doesnt change.
You just need to change how you calculate binomial coefficient, and that is all for this task.
Let just ignore the fact that though this need more detail, but as the blog is not about nCk problem I will just make it quick
For large prime $$$p > max(n, k)$$$
- Just using normal combinatorics related to factorial (since $$$p > max(n, k)$$$ nothing will affect the result)
- For taking divides under modulo you can just take modular inversion (as a prime always exist such number)
- Yet this is standard problem, just becareful of the overflow part
- You can also optimize by precalculating factorial, inversion number and inversion factorial in linear too
For general prime $$$p$$$
- We can just ignore factors $$$p$$$ in calculating $$$n!$$$.
- You also need to know how many times factor $$$p$$$ appears in $$$1 \dots n$$$
- Then combining it back when calculating for the answer.
- If we dont do this $$$n!$$$ become might divides some factors of $$$p$$$.
- By precalculation you can answer queries in $$$O(1)$$$
For squarefree $$$p$$$
- Factorize $$$p = p_1 \times p_2 \times p_q$$$ that all $$$p_i$$$ is prime.
- Ignore all factors $$$p_i$$$ when calculate $$$n!$$$.
- Remember to calculate how many times factors $$$p_i$$$ appear in $$$1 \dots n$$$.
- When query for the answer we just combine all those part back.
- Remember you can just take modulo upto $$$\phi(p)$$$ which you can also calculate while factorizing $$$p$$$.
- Remember that $$$n!$$$ must not divides any factor $$$p_i$$$ otherwise you will get wrong answer.
- By precalculation you can answer queries in $$$O(\log p)$$$
For general positive modulo $$$p$$$
- Factorize $$$p = p_1^{f_1} \times p_2^{f_2} \times p_q^{f_q}$$$ that all $$$p_i$$$ is unique prime.
- We calculate $$$C(n, k)$$$ modulo $$$p_i^{f_i}$$$ for each $$$i = 1 \dots q$$$.
- To do that, we need to calculate $$$n!$$$ modulo $$$p_i^{f_i}$$$ which is described here.
- To get the final answer we can use CRT.
- Yet this is kinda hard to code and debug also easy to make mistake so you must becareful
- I will let the implementation for you lovely readers.
- Yet depends on how you calculate stuffs that might increase your query complexity
- There are few (effective or atleast fully correct) papers about this but you can read the one written here
Implementation
O(n log mod + sqrt(mod)) for prime p or squarefree pvector<int> factor;
int factorize(int n) /// Calculating phi(n) while factorizing (n) in O(sqrt n)
{
factor.clear();
int phi = n;
if (!(n & 1))
{
n >>= __builtin_ctz(n);
factor.push_back(2);
phi -= phi / 2;
}
for (int x = 3; x * x <= n; x += 2)
{
if (n % x == 0)
{
do n /= x; while (n % x == 0);
factor.push_back(x);
phi -= phi / x;
}
}
if (n > 1)
{
factor.push_back(n);
phi -= phi / n;
}
return phi;
}
int f[LIM]; /// f[x] = nCk(n, x)
int fact[LIM]; /// n!
int tcaf[LIM]; /// n!^(-1)
int divp[LIM]; /// x but ignore all factors p[i]
int cntp[LIM][LOG_LIM]; /// cntp[x][i] = Number of time factor p[i] appear in 1..x
void precal(int MOD) /// Calculate f[x] for all x = 1 -> n in O(n log mod + sqrt mod)
{
int PHIMOD = factorize(MOD);
for (int x = 1; x <= n; ++x) /// For each part x in n!
{
int &t = divp[x] = x;
for (int i = 0; i < factor.size(); ++i) /// Ignore all factor p[i] of p
{
cntp[x][i] = cntp[x - 1][i];
for (; t % factor[i] == 0; t /= factor[i]) /// Count how many times p[i] appears in 1..n
++cntp[x][i];
}
}
fact[0] = fact[1] = 1;
tcaf[0] = tcaf[1] = 1;
for (int x = 2; x <= n; ++x) /// Finding n! and n!^(-1)
{
fact[x] = (1LL * fact[x - 1] * divp[x]) % MOD;
tcaf[x] = powMOD(fact[x], PHIMOD - 1, MOD);
}
memset(f, 0, sizeof(f[0]) * k);
for (int x = k; x <= n; ++x)
{
/// Calculate nCk % p normally
f[x] = fact[x];
mulMOD(f[x], tcaf[k], MOD);
mulMOD(f[x], tcaf[x - k], MOD);
for (int i = 0; i < factor.size(); ++i) /// Bringing those factors back
{
int p = cntp[x][i] - cntp[k][i] - cntp[x - k][i];
f[x] = 1LL * f[x] * powMOD(factor[i], p, MOD) % MOD;
}
}
}
Complexity
As you might notice $$$p$$$ is atleast prime, or atmost a squarefree number.
Since the complexity also depends on the factors of $$$p$$$, in the worst case, $$$p = 2 \times 3 \times 5 \times 7 \times 11 \times 13 \times 17 \times 19 \times \dots$$$ having most factors that is about $$$\log(p)$$$.
Also because of that factorizing and calculating $$$\phi{p}$$$ part take $$$O(\sqrt{p})$$$ time.
So you got $$$O(n \times \log p + \sqrt{p})$$$ in final.
Though you can still optimize this but by doing that why dont you just go straight up to solve for non squarefree $$$p$$$ too ?
Contribution
Yurushia for pointing out the linear complexity of squarefree sieve.
clyring for fixing typos, and the approach for tasks A, B, C, D, E, G, H, J