PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: sushil2006
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Easy
PREREQUISITES:
Dynamic programming
PROBLEM:
You’re given an array A consisting of ones and twos.
Define f(A) to be the number of prefix sums of A that are primes.
Let M = \max f(R) across all rearrangements R of A.
Find the minimum number of adjacent swaps in A needed to attain a rearrangement with value M.
EXPLANATION:
First, let’s find the value M, i.e. the maximum possible number of prime prefix sums.
Let S = A_1 + A_2 + \ldots + A_N denote the sum of the elements of A.
Clearly, any prefix sum must be \le S, so the absolute best we can do is to reach all primes that are \le S.
It turns out that this is (almost) always possible.
For example, one construction that attains this maximum is as follows:
- If every element of A equals 1, then N = S and we trivially reach the values 1, 2, 3, \ldots, N so the claim holds.
- If A contains both a 1 and a 2, we have the following construction:
[2, 1, 2, 2, \ldots, 2, 2, 1, 1, \ldots, 1]- That is, place a single 2, then a single 1, then all remaining 2's, and finally all remaining 1's.
- This obtains 2 and 3 as the initial prefix sums, then it goes through all odd numbers starting from 3 onwards.
Once the twos are used up, we place the remaining ones, going through all remaining numbers without skipping anything. - Because every prime other than 2 is odd, and we don’t skip any odd number \le S (except 1), every prime \le S will appear as a prefix sum.
The only exceptional case is when S contains only twos; in which case there’s only one rearrangement possible and the only prime it can reach is 2.
So, the single edge case aside, we never have to skip a prime number \le S = \text{sum}(A).
Now, we shift our focus to computing the minimum swaps needed.
To do that, the key observation is that because we have only ones and twos, we can focus on just one ‘type’ of value.
For example, let’s look at only the positions of the 1's.
Note that if these are fixed, then the positions of the 2's are automatically fixed too.
Further, note that if the 1's in A are initially at indices x_1, x_2, \ldots, x_k and finally at positions y_1, y_2, \ldots, y_k, then we need
swaps to achieve this configuration.
This method of representing the swap cost only in terms of the positions of the ones is quite helpful, and allows us to write a solution using dynamic programming.
Specifically, let’s define dp(i, j) to be the minimum cost such that:
- We’ve placed i ones and j twos in the prefix of length (i+j).
- Note that the sum of this prefix now equals i+2j.
- While doing so, we did not skip any prime number that’s \le i+2j.
Transitions are as follows.
The last element of the prefix can be either a one or a two.
First, consider the case that the last element in the prefix is a 1.
Then, we need to place i-1 ones and j twos in the prefix before this.
Further, this guarantees that the i-th one will end up at index (i+j), which as we saw above incurs a cost of |x_i - (i+j)| (given that its initial position is x_i).
Thus, the minimum cost of doing this is
Note that since we’re placing a 1, no value gets skipped, so we don’t need to worry about skipping a prime.
Next, consider the case that the last element in the prefix is a 1.
Note that this is only valid when (i+2j-1) is not a prime, since we’re skipping that value as a prefix sum.
In this case, the minimum cost of the previous prefix if dp(i, j-1), but that’s it - since we’ve rewritten the cost to depend purely on the positions of the 1's, placing a 2 here doesn’t affect the cost in any way (yet).
Putting the cases together,
- If i+2j-1 is not a prime, we have
dp(i, j) = \min(dp(i, j-1), dp(i-1, j) + |x_i - (i+j)|) - If i+2j-1 is a prime, we have dp(i, j) = dp(i-1, j) + |x_i - (i+j)|.
Needless to say, take care of the i = 0 and j = 0 cases appropriately (where it’s impossible to place any ones/twos respectively.)
The final answer is, of course, dp(c_1, c_2), where c_1, c_2 denote the counts of 1 and 2 in A respectively.
There are \mathcal{O}(N^2) states and \mathcal{O}(1) transitions from each one, so this is fast enough for the constraints.
TIME COMPLEXITY:
\mathcal{O}(N^2) 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());
const int N = 5005;
int dp[N][N];
int main()
{
ios::sync_with_stdio(false); cin.tie(0);
const int M = 10010;
vector prime(M, 0);
for (int i = 2; i < M; ++i) {
prime[i] = 1;
for (int j = 2; j < i; ++j) {
if (i%j == 0) {
prime[i] = 0;
break;
}
}
}
int t; cin >> t;
while (t--) {
int n; cin >> n;
vector a(n, 0);
for (int &x : a) cin >> x;
vector<int> ones;
for (int i = 0; i < n; ++i) {
if (a[i] == 1) ones.push_back(i);
}
if (size(ones) == 0) {
cout << 0 << '\n';
continue;
}
int p = size(ones);
int q = n - p;
for (int i = 0; i <= p; ++i) {
for (int j = 0; j <= q; ++j) {
dp[i][j] = 1e9;
if (i == 0 and j == 0) dp[i][j] = 0;
if (i > 0) {
int from = ones[i-1];
int to = i+j-1;
dp[i][j] = min(dp[i][j], dp[i-1][j] + abs(from - to));
}
if (j > 0 and !prime[i + 2*j - 1]) {
dp[i][j] = min(dp[i][j], dp[i][j-1]);
}
}
}
cout << dp[p][q] << '\n';
}
}