PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
TBD
PREREQUISITES:
Sieve, Set hashing
PROBLEM:
For a tree and an array A, f(A) is defined as follows:
- Edge (u, v) is given the value A_u + A_v.
- f(A) is the GCD of all edge values.
Given a tree and an array of length N-1, along with an integer M, for each 1 \leq x \leq M compute the minimum possible f(A) across all rearrangements of A assuming A_N = x.
EXPLANATION:
Read the solution to the easy version first, which solves the problem for a fixed array A.
Unfortunately, the Mobius DP solution that works there doesn’t directly extend to the harder version (to my knowledge).
Instead, let’s step back a bit.
We can still run the \mathcal{O}(N + M\log M) algorithm to compute all the ct(g, r) values, so let’s start with that.
Now, rather than fixing A_N and computing the optimal g, let’s fix g and see what that means for each A_N.
To begin with, if g falls under one of the impossible cases (more than two remainders, or the remainders are incompatible) then it can be ignored entirely.
If g has two compatible remainders, say r and g-r, then for any A_N that’s not congruent to either r or g-r this g does nothing.
So, there are only \mathcal{O}(\frac{M}{g}) important values of A_N, let’s try each of them and see what we get.
If A_N is fixed, either the count of r or g-r will increase by 1.
If the counts are (L, N-L) after this increase, we have the possibility of getting (a multiple of) g.
However, observe that there is, in a sense, a unique way of obtaining g like this.
Because the order of elements doesn’t really matter, we’re essentially dealing with a partition of the array into two sets - and there’s a unique way of forming this partition (well, technically two ways if L = N-L).
Let the subset of size L be S.
For now, we’ll store a map best[x][S], which gives the largest g for which, if we have A_N = x, the set S is a valid candidate.
Of course we can’t actually store S directly due to time/memory constraints but we’ll come back to this.
Next, suppose g has only one remainder r.
If r = 0 or r = \frac{g}{2}, A_N must also have the same remainder r to be valid.
For each of these A_N, no matter what is done, g will always divide all the edge values.
Let’s store the value force[x], which is the largest integer g which will always divide every edge value if we set A_N = x.
For a fixed g as above, \frac{M}{g} values of x will need their force[x] values updated.
What if r \neq 0 and r\neq \frac{g}{2}?
In this case, we can in fact just run the algorithm we did for two distinct remainders, allowing A_N to take the value g-r.
This is necessary when L = 1, for example.
Everything done above (other than actually storing best[x][S], since S can be large) takes \mathcal{O}(M\log M) time.
To optimize the set part, we use hashing: choose your favorite method of hashing sets (if done on indices) or multisets (if done on values), and store best[x][hash(S)] instead.
This way, we have \mathcal{O}(M\log M) map updates which is fast enough.
With all this precomputation done, let’s actually find the answer for some value x.
First, note that no matter what, the minimum GCD when A_N = x cannot fall below force[x].
Ideally, this would exactly be our answer - but is it always achievable?
It turns out, no, not necessarily.
In particular, let’s look at the best[x][hash(S)] values - or rather, the map best[x] itself.
Each value in best[x] corresponds to some partition that gives a GCD that’s greater than force[x].
So, to obtain a GCD of force[x], we’ll need to avoid all these partitions.
There are \binom{N}{L} partitions, so as long as best[x] has a size less than \binom{N}{L}, the answer is force[x].
If the size of best[x] equals \binom{N}{L}, we have no choice but to take one of them: in this case, the best answer is of course \min(best[x].values).
There is one minor detail: how exactly do we compute \binom{N}{L} and compare it against the size of best[x], when N can be large (so the coefficient is too large to store)?
To overcome this, note that we don’t need the exact value of \binom{N}{L}, and only want to know how it compares to size(best[x]).
In particular, if it grows too large we can just stop computing it immediately.
This allows the following algorithm:
- Ensure that L \leq N-L.
- Initialize the coefficient to 1.
- For each i from 1 to L, multiply the coefficient by N-i+1 and then divide it by i.
Essentially, this process computes \binom{N}{1}, \binom{N}{2}, \ldots, \binom{N}{L} in order.
Since we ensured L \leq N-L, this sequence is monotonic so breaking out when it exceeds the limit is valid.
TIME COMPLEXITY:
\mathcal{O}(N + M\log^2 M) per testcase.
CODE:
Editorialist's code (C++)
// #include <bits/allocator.h>
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());
int main()
{
ios::sync_with_stdio(false); cin.tie(0);
int t; cin >> t;
while (t--) {
int n, m; cin >> n >> m;
vector a(n-1, 0);
for (int &x : a) cin >> x;
vector adj(n, vector<int>());
for (int i = 0; i < n-1; ++i) {
int u, v; cin >> u >> v;
adj[--u].push_back(--v);
adj[v].push_back(u);
}
vector<int> col(n, 0);
{
vector<int> qu = {0};
col[0] = 1;
while (!qu.empty()) {
int u = qu.back(); qu.pop_back();
for (int v : adj[u]) {
if (col[v]) continue;
col[v] = 3 - col[u];
qu.push_back(v);
}
}
}
int L = ranges::count(col, 1);
int R = n-L;
if (L > R) swap(L, R);
vector<ll> mapped(n);
for (int i = 0; i < n; ++i) mapped[i] = uniform_int_distribution<ll>(0, (1ll << 62) - 1)(RNG);
vector freq(2*m + 1, 0);
vector here_hash(2*m + 1, 0ll);
for (int i = 0; i+1 < n; ++i) {
++freq[a[i]];
here_hash[a[i]] ^= mapped[i];
}
vector nxt(2*m + 2, 2*m + 1);
for (int i = 2*m; i >= 1; --i) {
nxt[i] = nxt[i+1];
if (freq[i]) nxt[i] = i;
}
vector<map<ll, int>> best(2*m + 1);
vector<int> forced(2*m + 1, 1);
for (int g = 1; g <= 2*m; ++g) {
int x = nxt[1];
vector<array<ll, 3>> values; // remainder, count, hash
while (x <= 2*m) {
bool have = false;
for (auto &[rem, ct, hsh] : values) {
if (rem == x%g) {
hsh ^= here_hash[x];
ct += freq[x];
have = true;
break;
}
}
if (!have) {
values.push_back({x%g, freq[x], here_hash[x]});
}
if (values.size() > 2) break;
x = nxt[x+1];
}
if (values.size() > 2) continue;
if (values.size() == 1) {
auto [rem, ct, hsh] = values[0];
if (rem != 0 and 2*rem != g) {
int want = g - rem;
for (int y = want; y <= 2*m; y += g) {
best[y][mapped[n-1]] = g;
if (n == 2) best[y][mapped[0]] = g;
}
continue;
}
for (int y = rem; y <= 2*m; y += g)
forced[y] = max(forced[y], g);
}
else {
auto [r1, c1, h1] = values[0];
auto [r2, c2, h2] = values[1];
if ((r1 + r2)%g != 0) continue;
for (int y = 0; y <= 2*m; y += g) {
int val = y + r1;
if (val <= m) {
if (c1 + 1 == L) best[val][h1 ^ mapped[n-1]] = g;
if (c2 == L) best[val][h2] = g;
}
val = y + r2;
if (val <= m) {
if (c1 == L) best[val][h1] = g;
if (c2 + 1 == L) best[val][h2 ^ mapped[n-1]] = g;
}
}
}
}
ll choices = 1;
for (int i = 1; i <= L; ++i) {
choices *= n-i+1;
choices /= i;
if (choices > 1e9) break;
}
for (int x = 1; x <= m; ++x) {
if (size(best[x]) == choices) {
// Every subset forces something
int ans = 2*m + 10;
for (auto [sub, g] : best[x]) ans = min(ans, g);
cout << ans << ' ';
}
else {
cout << forced[x] << ' ';
}
}
cout << '\n';
}
}