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.
Notice that in this blog, we will solve in $$$O(n)$$$ for general $$$k$$$ from the original task
For harder variants you can see in this blog [Variants] An interesting counting problem related to square product
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;
}
Extra Tasks
These are harder variants, and generalization from the original problem. You can see more detail here
*Marked as solved only if tested with atleast $$$10^6$$$ queries
A: Can we also use phi function or something similar to solve for $$$k = 3$$$ in $$$O(\sqrt{n})$$$ or faster ?
B: Can we also use phi function or something similar to solve for general $$$k$$$ in $$$O(\sqrt{n})$$$ or faster ?
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)$$$ ?
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 also solve the problem where there are no order: Just simply $$$0 \leq a_i \leq n$$$ ?
M: 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$$$) ?
N: 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$$$ ?