PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Dynamic programming
PROBLEM:
You’re given an array A. At most K times, you can increment a single element of A.
Find the maximum possible value of \sum \text{sep}_A(x), where:
- If x\not\in A then \text{sep}_A(x) = 0.
- Otherwise \text{sep}_A(x) equals the largest distance between any two occurrences of x in A.
EXPLANATION:
Clearly, only the endpoints of each value matter to compute its separation.
Thus, one interpretation of what we’re trying to do, is to choose several tuples (l, r, x) such that x \ge \max(A_l, A_r) for each tuple (meaning we’ll make A_l and A_r both equal to x eventually.)
This tuple adds r-l to our score, and has a cost of 2x-A_l-A_r.
Further, there must be at most one tuple corresponding to each value x, and each index in [1, N] can appear in at most one tuple as well.
Now, observe that the actual (l, r, x) tuples don’t matter so much as which indices are chosen as left/right endpoints in some tuple.
That is, it doesn’t really matter if we have (l_1, r_1, x_1) and (l_2, r_2, x_2) or (l_1, r_2, x_1) and (l_2, r_1, x_2), since the overall score and cost remain the same either way - of course as long as the x \ge \max(A_l, A_r) constraint is maintained.
So, we can focus on just deciding which indices are left endpoints and which ones are right endpoints.
To deal with the x \ge \max(A_l, A_r) condition, we’ll process indices in ascending value of their elements.
So, we’ll process x = 1, 2, 3, \ldots, N in order; and whenever we’re processing x, all indices containing a copy of x will become available.
Now, when an index i is activated, there are three options for what to do with it: do nothing, or use it as a left/right endpoint.
Using it is a left/right endpoint adds -i or +i to the score, respectively, and a constant cost of -A_i.
However, the cost of using this index also depends on which value x it finally satisfies - after all, if the final value is y then it needed (y - A_i) increments.
However, one way to think about this is that: for each index that we use, as long as it is active, its cost will increase by 1 each time we increase x.
This can be done simultaneously for all active elements - and so we only really need to know the number of active elements at any point of time in order to update the cost.
With this, we have a fairly straightforward (if slow) DP available to us.
Upto the current value x, define dp[L][R][c] to be the maximum possible score such that there are L alive left endpoints, R alive right endpoints, and the total cost so far is c.
When processing an index with A_i = x, the transitions are simple:
- Do nothing, or
- Make this a left endpoint: so L increases by 1, and the score decreases by i.
- Make this a right endpoint: so R increases by 1, and the score increases by i.
After processing all indices with value x, we’ll then update this DP a bit more as well.
First, observe that if we have both L \gt 0 and R \gt 0 at some state, it’s optimal to immediately pair a left endpoint with a right one.
This is because keeping them alive doesn’t increase their score, but does increase cost.
So, such a state can update its value to dp[L-1][R-1][c] instead.
Finally, when moving from x to x+1, the cost corresponding to all existing active indices increases by 1.
So, each dp[L][R][c] increases by L+R.
The above DP is indeed correct, but too slow - there are \mathcal{O}(N^2 K) states, and each element of the array along with each transition x\to x+1 updates all of them for \mathcal{O}(N^3 K) in total.
For the constraints, this is too slow.
The key observation to speed this up is that L and R, the number of active endpoints, can’t actually be too big.
Indeed, suppose we have L active left endpoints currently at x.
Then, we can use up at most one of them at x, then at most one more at x+1, and so on till we use up the last one no earlier than x+L-1.
The issue is, this costs too much!
From these left endpoints alone, the x\to x+1 transition has a cost of at minimum (L-1) + (L-2) + \ldots + 1 = \frac{L\cdot (L-1)}{2}.
We need this to be \le K, which is only possible for approximately L \le \sqrt{2K}.
In practice, for K = 400, L is bounded above by 27.
The exact same analysis applies to R, and so R is effectively bounded above by 27 as well.
Observe that this simple observation reduces the number of states in the DP from \mathcal{O}(N^2 K) to \mathcal{O}(\sqrt K \cdot \sqrt K \cdot K) = \mathcal{O}(K^2).
Accounting for transitions, the overall complexity is \mathcal{O}(NK^2) which is fast enough for the constraints.
TIME COMPLEXITY:
\mathcal{O}(NK^2) per testcase.
CODE:
Tester's code (C++)
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
template<typename T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long long int ll;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL)
#define pb push_back
#define endl '\n'
#define sz(a) (int)a.size()
#define setbits(x) __builtin_popcountll(x)
#define ff first
#define ss second
#define conts continue
#define ceil2(x,y) ((x+y-1)/(y))
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define yes cout << "Yes" << endl
#define no cout << "No" << endl
#define rep(i,n) for(int i = 0; i < n; ++i)
#define rep1(i,n) for(int i = 1; i <= n; ++i)
#define rev(i,s,e) for(int i = s; i >= e; --i)
#define trav(i,a) for(auto &i : a)
template<typename T>
void amin(T &a, T b) {
a = min(a,b);
}
template<typename T>
void amax(T &a, T b) {
a = max(a,b);
}
#ifdef LOCAL
#include "debug.h"
#else
#define debug(...) 42
#endif
/*
*/
const int MOD = 1e9 + 7;
const int N = 1e5 + 5;
const int inf1 = int(1e9) + 5;
const ll inf2 = ll(1e18) + 5;
void solve(int test_case){
ll n,k; cin >> n >> k;
vector<ll> a(n+5);
rep1(i,n) cin >> a[i];
vector<ll> pos[n+5];
rep1(i,n) pos[a[i]].pb(i);
ll s = 0;
ll siz = -1;
rep1(i,k+1){
s += i;
if(s > k){
siz = i;
break;
}
}
// siz = max #of + or max #of - we can have at any point of time
ll dp[siz+5][siz+5][k+5]; // max cost for given (+,-,cost)
memset(dp,-0x3f,sizeof dp);
dp[0][0][0] = 0;
rep1(x,n){
// introduce any new guys at given x value
trav(i,pos[x]){
rev(plus_cnt,siz,0){
rev(minus_cnt,siz,0){
rep(cost,k+1){
ll curr = dp[plus_cnt][minus_cnt][cost];
// make +
amax(dp[plus_cnt+1][minus_cnt][cost], curr+i);
// make -
amax(dp[plus_cnt][minus_cnt+1][cost], curr-i);
}
}
}
}
// transition x --> x+1
rev(cost,k,0){
rev(plus_cnt,siz,0){
rev(minus_cnt,siz,0){
ll curr = dp[plus_cnt][minus_cnt][cost];
ll sub = plus_cnt > 0 and minus_cnt > 0;
ll plus_cnt_new = plus_cnt-sub;
ll minus_cnt_new = minus_cnt-sub;
ll new_cost = cost+plus_cnt_new+minus_cnt_new;
if(new_cost <= k){
amax(dp[plus_cnt_new][minus_cnt_new][new_cost], curr);
}
if(!(plus_cnt == 0 and minus_cnt == 0)){
dp[plus_cnt][minus_cnt][cost] = -inf2;
}
}
}
}
}
ll ans = 0;
rep(plus_cnt,siz+1){
ll minus_cnt = plus_cnt;
rep(cost,k+1){
ll curr = dp[plus_cnt][minus_cnt][cost];
ll to_add = plus_cnt*(plus_cnt-1)/2;
ll new_cost = cost + 2*to_add;
if(new_cost <= k){
amax(ans,curr);
}
}
}
cout << ans << endl;
}
int main()
{
fastio;
int t = 1;
cin >> t;
rep1(i, t) {
solve(i);
}
cerr << "RUN SUCCESSFUL" << endl;
return 0;
}