For the rest of the problem, define $$$f_i(x)$$$ as the number of people up to person $$$i$$$ (zero-indexed) with rank $$$x$$$. For convenience, let $$$f_{-1}(x) = 0$$$ for all $$$x$$$.
We shall start with part (a). The first observation is as follows: If person $$$i$$$ with country rank $$$x$$$ has some person $$$j$$$ in front of them who must be in the same country, then there must be some person $$$i$$$ with country rank $$$x-1$$$ in front of them who must be in the same country.
A quick idea of the proof would be that one can swap the country of person $$$i$$$ and some other people in the same country to change person $$$j$$$. For the formal proof, suppose otherwise. Then there must be at least 2 possibilities of persons $$$a$$$ and $$$b$$$, both with country rank $$$x-1$$$. Suppose without loss of generality person $$$i$$$ is in the same country as person $$$a$$$ in some construction. Then consider swapping the countries of everyone in person $$$a$$$'s country with a worse rank that $$$x-1$$$ and everyone in person $$$b$$$'s country with a worse rank than $$$x-1$$$. One still gets a valid configuration of countries, however person $$$i$$$ has now switched countries and has a different person $$$j$$$ in front of them, a contradiction since person $$$j$$$ is determined.
To continue, consider processing person $$$i$$$ with country rank $$$x$$$ and determining which people in front of them must be from the same country as them. From the observation, it is sufficient to determine the person with rank $$$x-1$$$ (if there exists any) since we can mark them as the same country and in the end, count the number of people connected together (this calls for DSU / UFDS to keep track of who must be together). Consider $$$f_{i-1}(x-1) - f_{i-1}(x)$$$, which is the number of people with country rank $$$x-1$$$ which have not been picked by someone with a country rank $$$x$$$. It is guaranteed to be at least 1 since a valid ranking must exist. If $$$f_{i-1}(x-1) - f_{i-1}(x) \gt 1$$$, there are at least 2 possible people with country rank $$$x-1$$$ who are in the same country. Thus $$$f_{i-1}(x-1) - f_{i-1}(x) = 1$$$ is necessary.
It is, however, not sufficient. Consider a case where the country ranks are $$$[0, 0, 1, 1]$$$. For the last person, $$$f_2(0)-f_2(1)=2-1=1$$$, however it is still possible for them to be in the same country as either one of the 0s. After some thinking and/or guessing, we obtain that $$$f_{i-1}(x-1) - f_{i-1}(x) = 1$$$ must hold, and the previous occurrence of country rank $$$x-1$$$ must be closer than the previous occurrence of country rank $$$x$$$. Another way of looking at this is that if $$$f_{i-1}(x-1) - f_{i-1}(x)$$$ ever exceeds 1, anyone following with country rank $$$x$$$ will not have a guaranteed partner until $$$f_{i-1}(x-1) - f_{i-1}(x)$$$ "resets" to 0.
As mentioned earlier, we can use a UFDS / DSU to maintain relationships of who is in the same country. In my code prog denotes $$$f_{i-1}(x) - f_{i-1}(x+1)$$$ while rs denotes whether the values of prog have been reset to 0. There is the special case of someone with country rank $$$0$$$.
#include "country.h"
#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define iloop(m, h) for (auto i = m; i != h; i += (m<h?1:-1))
#define jloop(m, h) for (auto j = m; j != h; j += (m<h?1:-1))
#define kloop(m, h) for (auto k = m; k != h; k += (m<h?1:-1))
#define pll pair<ll, ll>
ll par[1000005], sz[1000005];
ll n, cans, tmp;
ll find(ll x) {
return (x == par[x] ? x : par[x] = find(par[x]));
}
void merge(ll x, ll y) {
x = find(x), y = find(y);
if (x == y) return;
if (sz[x] < sz[y]) swap(x, y);
par[y] = x;
sz[x] += sz[y];
}
long long count_same_country(int N, vector<int> country_rank) {
n = N, cans = 0;
ll prog[n], a[n], prv[n];
bool rs[n];
iloop(0, n) a[i] = country_rank[i], par[i] = i, sz[i] = 1, prog[i] = 0, rs[i] = 1;
iloop(0, n) {
if (a[i]) {
if (prog[a[i]-1] == 1 && rs[a[i]-1]) merge(i, prv[a[i]-1]);
prog[a[i]-1]--;
if (prog[a[i]-1] == 0) rs[a[i]-1] = 1;
}
prog[a[i]]++;
if (prog[a[i]] == 2) rs[a[i]] = 0;
prv[a[i]] = i;
}
iloop(0, n) if (par[i] == i) cans += sz[i] * (sz[i] - 1) / 2;
return cans;
}
For part (b), it helps to consider when 2 people cannot be in the same country. There are 3 possibilities:
- The person in front has a country rank at least that of the person behind.
- There are not enough people in between that can form the chain of people in the same country with successive country ranks.
- The person in front is forced to be connected to someone else instead of that person, like in the scenario $$$[0, 1, 0, 1]$$$ where the last person cannot be with the first person.
The first 2 cases seem reasonable to handle. The first case contributes a lower bound to what country ranks a new person must be to be in the same country, while the second case seems to contribute an upper bound. It turns out that the third case also contributes a lower bound!
The third case borrows a similar idea from part (a): When $$$f_{i}(x-1) - f_{i}(x) = 0$$$, everyone up to that point with country rank $$$x-1$$$ is forced to have the next person from the same country to be indexed at most $$$i$$$. Usefully, $$$f_i(x)$$$ is monotone decreasing on $$$x$$$.
It can be shown (with many details, unfortunately not by me) that these conditions are necessary and sufficient for 2 people to be forced into different countries. In contest, I first implemented a $$$O(N^2)$$$ solution to verify the claim (it gave 21 points on (b)). Processing everyone in increasing order of rank, I assign everyone a range $$$[lh, rh]$$$ where a person can potentially share a country with someone new being processed if the country rank of the new person lies in the range. For the 3 possibilities,
- At the start
lh = rh = country_rank[i]+1 - If some person with country rank
rh shows up, increase rh by 1 - If the count of $$$f_i{lh} - f_i{lh+1}$$$ reaches 0, increment $$$lh$$$
This is easier to check from the perspective of the new person being processed:
- For anyone with
rh as the current country rank, increment it - If the increase in country rank results in $$$f_i{lh} - f_i{lh+1} = 0$$$ where
lh + 1 = country_rank[i], increment all such $$$lh$$$
Then for each person, the number of people in front who cannot be in the same country is the number of people whose range does not include the country rank.
The code for this partial solution is as follows:
long long count_diff_country(int N, vector<int> country_rank) {
ll n = N, cans = 0;
ll prog[n], a[n], lh[n], rh[n];
iloop(0, n) a[i] = country_rank[i], prog[i] = 0;
iloop(0, n) {
jloop(0, i) {
if (lh[j] > a[i] || rh[j] < a[i]) cans++;
if (rh[j] == a[i]) rh[j]++;
}
if (a[i]) {
prog[a[i]-1]--;
if (!prog[a[i]-1]) jloop(0, i) if (lh[j] == a[i]) lh[j]++;
}
prog[a[i]]++;
lh[i] = rh[i] = a[i]+1;
}
return cans;
}
To speed this up, notice that we only care about the values of $$$lh$$$ and $$$rh$$$ and not about the indices that correspond to them. We can thus use a data structure such as a segment tree to store these values. The AC code for this part is as follows:
ll seg[2][2000005];
void upd(ll X, ll V, ll id) {
for (seg[id][X += n] += V; X > 1; X >>= 1) seg[id][X>>1] = seg[id][X] + seg[id][X^1];
}
ll qu(ll S, ll E, ll id) {
E++;
ll res = 0;
for (S += n, E += n; S != E; S >>= 1, E >>= 1) {
if (S&1) res += seg[id][S++];
if (E&1) res += seg[id][--E];
}
return res;
}
long long count_diff_country(int N, vector<int> country_rank) {
n = N, cans = 0;
ll prog[n], a[n];
iloop(0, 2*n) seg[0][i] = seg[1][i] = 0;
iloop(0, n) a[i] = country_rank[i], prog[i] = 0;
iloop(0, n) {
cans += qu(a[i]+1, n-1, 0);
cans += qu(0, a[i]-1, 1);
tmp = qu(a[i], a[i], 1);
upd(a[i], -tmp, 1);
upd(a[i]+1, tmp, 1);
if (a[i]) {
prog[a[i]-1]--;
//if (!prog[a[i]-1]) jloop(0, i) if (lh[j] == a[i]) lh[j]++;
if (!prog[a[i]-1]) {
tmp = qu(a[i], a[i], 0);
upd(a[i], -tmp, 0);
upd(a[i]+1, tmp, 0);
}
}
prog[a[i]]++;
upd(a[i]+1, 1, 0);
upd(a[i]+1, 1, 1);
}
return cans;
}
You may try to see how the 2 pieces of code here basically do the same thing. This now runs in $$$O(N\logN)$$$