PROBLEM LINK:
Practice
Div-4 Contest
Div-3 Contest
Div-2 Contest
Div-1 Contest
Author: Vladislav Bargatin, Tejas Pandey
Tester: Aryan Choudhary, Istvan Nagy
Editorialist: Daanish Mahajan
DIFFICULTY:
Easy Medium
PREREQUISITES:
Greedy
PROBLEM:
Chef found an array A of N elements. He defines a subarray as bad if the maximum element of the subarray is equal to the minimum element of the subarray.
Chef is allowed to change at most K elements of the array.
Find the minimum number of bad subarrays Chef can obtain after changing at most K elements of the array.
EXPLANATION:
Hint 1
max(A_L, A_{L+1}, \ldots , A_R) = min(A_L, A_{L+1}, \ldots , A_R) means elements in the range [L, R] are same.
Hint 2
If we visualize changing the value of an element as partitioning the array about that index and a and b are the sizes of the two partitions, then in an optimal partition (the one with minimum bad subarrays) |a - b| \le 1.
Hint 3
Let f(m, k) denote the number of subsegments that satisfy the equality condition in a segment of length m after making k edits, then:
f(m, k + 1) - f(m, k) \le f(m, k + 2) - f(m, k + 1).
Solution
max(A_L, A_{L+1}, \ldots , A_R) = min(A_L, A_{L+1}, \ldots , A_R) means elements in the range [L, R] are same.
Proof
It’s easy to prove it by contradiction and left as an assignment for readers.
Starting from index i = 1 we find the closest index(j) such that all the elements in the range [i, j] are same. Next we repeat the process from index j + 1 and so on.
This results in the array getting divided into continuous non-zero length segments such that every element of the array belongs to exactly one segment.
So the problem is reduced to independent segments of length suppose L_1, L_2, \ldots, L_x where K_1, K_2, \ldots, K_x number of elements from each segment are changed ensuring \sum_{i = 1}^{x} K_i = K and 0 \le K_i \le L_i.
For optimal choice of K_i elements in a particular segment S_i, visualise that the elements chosen partition the segment into K_i + 1 smaller segments of lengths L_{i, 1}, L_{i, 2}, \ldots L_{i, K_i + 1}.
So the contribution reduces from:
Contribution of a big segment of length L_i \rightarrow Contribution of K_i + 1 small segments + Contribution of K_i unit length segments.
\frac{L_i \cdot (L_i + 1)}{2} \rightarrow \sum_{j = 1}^{K_i + 1} \frac{L_{i, j} \cdot (L_{i, j} + 1)}{2} + K_i
Right side is lesser than left side
Start with a segment being divided into two parts of lengths a and b respectively.
So expression reduces from:
\frac{(a + b + 1) \cdot (a + b + 2)}{2} \rightarrow \frac{a \cdot (a + 1)}{2} + \frac{b \cdot (b + 1)}{2} + 1
2 \cdot (a + b + a \cdot b) \rightarrow 0
This means with each step answer reduces.
So the right choice of elements to be modified is crucial. Observe that if there are any two partitions j_1, j_2 such that |L_{i, j_1} - L_{i, j_2}| \ge 2, it’s always better to increase the length of the smaller partition and reduce the length of the larger partition by 1 unit and keep doing it until there are no two partitions following this condition.
Proof
Consider two partitions of lengths a and a + x where x \ge 2. Now we update these partitions to a + 1 and a + x - 1. So contribution changes from:
\frac{a \cdot (a + 1)}{2} + \frac{(a + x) \cdot (a + x + 1)}{2} \rightarrow \frac{(a + 1) \cdot (a + 2)}{2} + \frac{(a + x - 1) \cdot (a + x)}{2}
which implies 2 \cdot x \rightarrow 2, therefore contribution reduces.
In the end, any two partitions will vary by at max 1 unit.
Let f(m, k) denote the number of subsegments that satisfy the equality condition in a segment of length m after making k edits. Then under optimal splitting:
Let p = (m - k) \bmod (k + 1), q = k + 1 - p, l = \frac{m - k}{k + 1}, then:
f(m, k) = \frac{q \cdot l \cdot (l + 1)}{2} + \frac{p \cdot (l + 1) \cdot (l + 2)}{2} + k
Now the last but most crucial observation is:
f(m, k + 1) - f(m, k) \le f(m, k + 2) - f(m, k + 1).
Proof
You can check it using brute force for all possible values of m and k locally.
This observation allows us to use greedy approach in which we add pair<long long, pair<int, int>> \{f(m, k + 1) - f(m, k), \{m, k \} \} for all segments to a multiset or priority queue. Here m denotes the length of the segment and k denotes the number of elements that have been changed in that segment. Refer to the linked solutions for implementation details.
COMPLEXITY ANALYSIS:
\mathcal {O}(N + (K + x) \log x).
TESTER’s NOTES:
Problem - F - Codeforces has MINMAXEQ as subproblem. That’s why up-solving is important.
"For me, 1661F on paper was more or less 5 mins solve because I have already analyzed the cost function in MINMAXEQ"… aryanc403 _/\_.
SOLUTIONS:
Setter's Solution
#include <bits/stdc++.h>
#define ll long long int
using namespace std;
ll sq(ll n) {
return (n*(n + 1))/2;
}
ll calculate(ll n, ll k) {
if(k >= n/2) return n;
k++;
ll first = (n + 1 - k)%k;
ll second = k - first;
ll sz = (n + 1 - k)/k;
return (k - 1 + first*sq(sz + 1) + second*sq(sz));
}
int main() {
//freopen("inp11.in", "r", stdin);
//freopen("out11.out", "w", stdout);
int t;
cin >> t;
assert(1 <= t && t <= 30000);
int sn = 0;
while(t--) {
int n, k;
cin >> n >> k;
sn += n;
assert(sn <= 300000);
assert(1 <= n && n <= 100000);
assert(1 <= k && k <= n);
int a[n];
for(int i = 0; i < n; i++) { cin >> a[i];
assert(1 <= a[i] && a[i] <= 1000000000);
}
ll ans = 0;
priority_queue<tuple<ll, ll, ll>> pq;
int now = 1;
for(int i = 1; i < n; i++) {
if(a[i] == a[i - 1]) now++;
else {
ans += calculate(now, 0);
pq.push({calculate(now, 0) - calculate(now, 1), now, 0});
now = 1;
}
}
ans += calculate(now, 0);
pq.push({calculate(now, 0) - calculate(now, 1), now, 0});
while(k) {
k--;
tuple<ll,ll,ll> res = pq.top();
pq.pop();
ans -= get<0>(res);
ll sz = get<1>(res), divide = get<2>(res);
divide++;
pq.push({calculate(sz, divide) - calculate(sz, divide + 1), sz, divide});
}
cout << ans << "\n";
}
}
Tester's Solution
/* in the name of Anton */
/*
Compete against Yourself.
Author - Aryan (@aryanc403)
Atcoder library - https://atcoder.github.io/ac-library/production/document_en/
*/
#ifdef ARYANC403
#include <header.h>
#else
#pragma GCC optimize ("Ofast")
#pragma GCC target ("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx")
//#pragma GCC optimize ("-ffloat-store")
#include<bits/stdc++.h>
#define dbg(args...) 42;
#endif
// y_combinator from @neal template https://codeforces.com/contest/1553/submission/123849801
// http://www.open-std.org/jtc1/sc22/wg21/docs/papers/2016/p0200r0.html
template<class Fun> class y_combinator_result {
Fun fun_;
public:
template<class T> explicit y_combinator_result(T &&fun): fun_(std::forward<T>(fun)) {}
template<class ...Args> decltype(auto) operator()(Args &&...args) { return fun_(std::ref(*this), std::forward<Args>(args)...); }
};
template<class Fun> decltype(auto) y_combinator(Fun &&fun) { return y_combinator_result<std::decay_t<Fun>>(std::forward<Fun>(fun)); }
using namespace std;
#define fo(i,n) for(i=0;i<(n);++i)
#define repA(i,j,n) for(i=(j);i<=(n);++i)
#define repD(i,j,n) for(i=(j);i>=(n);--i)
#define all(x) begin(x), end(x)
#define sz(x) ((lli)(x).size())
#define pb push_back
#define mp make_pair
#define X first
#define Y second
#define endl "\n"
typedef long long int lli;
typedef long double mytype;
typedef pair<lli,lli> ii;
typedef vector<ii> vii;
typedef vector<lli> vi;
const auto start_time = std::chrono::high_resolution_clock::now();
void aryanc403()
{
#ifdef ARYANC403
auto end_time = std::chrono::high_resolution_clock::now();
std::chrono::duration<double> diff = end_time-start_time;
cerr<<"Time Taken : "<<diff.count()<<"\n";
#endif
}
long long readInt(long long l, long long r, char endd) {
long long x=0;
int cnt=0;
int fi=-1;
bool is_neg=false;
while(true) {
char g=getchar();
if(g=='-') {
assert(fi==-1);
is_neg=true;
continue;
}
if('0'<=g&&g<='9') {
x*=10;
x+=g-'0';
if(cnt==0) {
fi=g-'0';
}
cnt++;
assert(fi!=0 || cnt==1);
assert(fi!=0 || is_neg==false);
assert(!(cnt>19 || ( cnt==19 && fi>1) ));
} else if(g==endd) {
if(is_neg) {
x=-x;
}
assert(l<=x&&x<=r);
return x;
} else {
assert(false);
}
}
}
string readString(int l, int r, char endd) {
string ret="";
int cnt=0;
while(true) {
char g=getchar();
assert(g!=-1);
if(g==endd) {
break;
}
cnt++;
ret+=g;
}
assert(l<=cnt&&cnt<=r);
return ret;
}
long long readIntSp(long long l, long long r) {
return readInt(l,r,' ');
}
long long readIntLn(long long l, long long r) {
return readInt(l,r,'\n');
}
string readStringLn(int l, int r) {
return readString(l,r,'\n');
}
string readStringSp(int l, int r) {
return readString(l,r,' ');
}
void readEOF(){
assert(getchar()==EOF);
}
vi readVectorInt(int n,lli l,lli r){
vi a(n);
for(int i=0;i<n-1;++i)
a[i]=readIntSp(l,r);
a[n-1]=readIntLn(l,r);
return a;
}
bool isBinaryString(const string s){
for(auto x:s){
if('0'<=x&&x<='1')
continue;
return false;
}
return true;
}
// #include<atcoder/dsu>
// vector<vi> readTree(const int n){
// vector<vi> e(n);
// atcoder::dsu d(n);
// for(lli i=1;i<n;++i){
// const lli u=readIntSp(1,n)-1;
// const lli v=readIntLn(1,n)-1;
// e[u].pb(v);
// e[v].pb(u);
// d.merge(u,v);
// }
// assert(d.size(0)==n);
// return e;
// }
const lli INF = 0xFFFFFFFFFFFFFFFL;
lli seed;
mt19937 rng(seed=chrono::steady_clock::now().time_since_epoch().count());
inline lli rnd(lli l=0,lli r=INF)
{return uniform_int_distribution<lli>(l,r)(rng);}
class CMP
{public:
bool operator()(ii a , ii b) //For min priority_queue .
{ return ! ( a.X < b.X || ( a.X==b.X && a.Y <= b.Y )); }};
void add( map<lli,lli> &m, lli x,lli cnt=1)
{
auto jt=m.find(x);
if(jt==m.end()) m.insert({x,cnt});
else jt->Y+=cnt;
}
void del( map<lli,lli> &m, lli x,lli cnt=1)
{
auto jt=m.find(x);
if(jt->Y<=cnt) m.erase(jt);
else jt->Y-=cnt;
}
bool cmp(const ii &a,const ii &b)
{
return a.X<b.X||(a.X==b.X&&a.Y<b.Y);
}
lli CostLen(const lli n){
return n*(n+1)/2;
}
lli calculate(lli n,lli k){
if(n<=2*k+1) return n;
lli ans=k;
n-=k;k++;
const lli len=n/k,ext=n%k;
ans+=ext*CostLen(len+1);
ans+=(k-ext)*CostLen(len);
return ans;
}
int main(void) {
ios_base::sync_with_stdio(false);cin.tie(NULL);
// freopen("txt.in", "r", stdin);
// freopen("txt.out", "w", stdout);
// cout<<std::fixed<<std::setprecision(35);
lli T,n,k;
T=readIntLn(1,3e4);
lli sumN = 3e5;
while(T--)
{
n=readIntSp(1,min(100000LL,sumN));
sumN-=n;
k=readIntLn(1,n);
auto a=readVectorInt(n,1,1e9);
priority_queue<pair<lli,ii>> pq;
lli pvr=a[0],cnt=0;
lli ans=0;
a.pb(-1);
for(auto x:a){
if(pvr==x)
cnt++;
else{
pq.push({calculate(cnt,0)-calculate(cnt,1),{cnt,0}});
ans+=CostLen(cnt);
cnt=1;
}
pvr=x;
}
while(k--){
const lli len=pq.top().Y.X,pieces=pq.top().Y.Y;
pq.pop();
ans+=calculate(len,pieces+1)-calculate(len,pieces);
pq.push({calculate(len,pieces+1)-calculate(len,pieces+2),{len,pieces+1}});
}
cout<<ans<<endl;
} aryanc403();
readEOF();
return 0;
}
Editorialist's Solution
#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define pll pair<ll, ll>
#define plll pair<ll, pll>
#define pb push_back
#define mp make_pair
#define F first
#define S second
multiset<plll> s;
ll cost(int x, int p){
ll rem = (x - p) % (p + 1), val = (x - p) / (p + 1);
return p + (p + 1 - rem) * val * (val + 1) / 2 + rem * (val + 1) * (val + 2) / 2;
}
int main() {
ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);
int t; cin >> t;
int n, k;
while(t--){
s.clear();
cin >> n >> k;
int pv = -1;
ll ans = 0, cnt = 0;
int x;
for(int i = 1; i <= n; i++){
cin >> x;
if(x != pv){
if(cnt != 0){
ans += cnt * (cnt + 1) / 2;
s.insert({cost(cnt, 1) - cost(cnt, 0), {cnt, 1}});
}
cnt = 1;
}else{
cnt++;
}
pv = x;
}
ans += cnt * (cnt + 1) / 2;
s.insert({cost(cnt, 1) - cost(cnt, 0), {cnt, 1}});
while(s.size() > 0 && k--){
plll p = *s.begin();
s.erase(s.begin());
// cout << p.F << endl;
ans += p.F;
if(p.S.F >= p.S.S + 1)s.insert({cost(p.S.F, p.S.S + 1) - cost(p.S.F, p.S.S), {p.S.F, p.S.S + 1}});
}
cout << ans << endl;
}
return 0;
}