PROBLEM LINK:
Setter: Ildar Gainullin
Tester: Alexander Morozov
Editorialist: Ajit Sharma Kasturi
DIFFICULTY:
EASY
PREREQUISITES:
None
PROBLEM:
We are given two integers N and K where K \leq N, we have to construct an array of N integers where the element A_i is either i or -i and there exists exactly K values of i such that
1 \leq i<N and A_1 + A_2 + ... +A_i>0 .
QUICK EXPLANATION:
One of the possible arrays satisfying the above conditions can be constructed as follows:
-
Case 1: K \leq N/2
The final array will be of the form
A_i = \begin{cases} i & \text{if } i \leq 2 \cdot K \ \text{and } i \ \text{is } \text{odd} ,\\ -i & \text{otherwise } \end{cases} -
Case 2: K>N/2
The final array will be of the form
A_i = \begin{cases} -i & \text{if } i \leq 2 \cdot (N-K) \ \text{and } i \ \text{is } \text{even} ,\\ i & \text{otherwise } \end{cases}
EXPLANATION:
There are various ways to construct the array to satisfy the final requirements. Let us first start initially with the following array:
- For 1 \leq i \leq N
A_i = \begin{cases} i & \text{if } i \ \text{is } \text{odd} ,\\ -i & \text{otherwise } \end{cases}
Let us take an example. If N is 5, then we have the array as follows:
A=[1, -2, 3 , -4, 5]
There are some interesting properties of this array.
-
Property 1: If i is even then A_1+A_2+...+A_i<0 .
Proof: A_1+A_2+...+A_i = (A_1+A_2)+(A_3+A_4)+...+(A_{i-1}+A_i)
\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ = (1-2)+(3-4)+...+((i-1)-i)
\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ = (-1)+(-1)+...+(-1)
\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ = \dfrac{-i}{2} <0
-
Property 2: If i is odd then A_1+A_2+...+A_i>0 .
Proof: A_1 =1 >0 . Now for i>1
\ \ \ \ \ \ \ \ \ \ A_1+A_2+...+A_i = (A_1+A_2+...+A_{i-1}) + A_i
\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ = \dfrac{-(i-1)}{2} +i \ \ \text{(from the proof of previous property)}
\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ = \dfrac{(i+1)}{2} >0
Therefore, after constructing this array, we will have \dfrac{(N+1)}{2} values of i for which A_1 + A_2 +... +A_i>0 since there will be \dfrac{(N+1)}{2} odd indices in the array . But what we want is exactly K values of i at which A_1 + A_2 +... +A_i>0 . So for that we can do the following:
-
If K> \dfrac{(N+1)}{2}
It is obvious that at exactly (K-\dfrac{(N+1)}{2}) values of i for which A_i=-i , we need to change A_i to i. So the intuitive way of doing that is to keep updating the values at those indices starting from the end of the array so that in each step after updating at some index i, the values of prefix sums for indices from 1 to i-1 won’t get affected and in each step we are actually increasing the possible values of i for which A_1 +A_2 +...+ A_i>0 by 1. -
If K < \dfrac{(N+1)}{2}
It is obvious that at exactly (\dfrac{(N+1)}{2}-K) values of i for which A_i=i , we need to change A_i to -i. So the intuitive way of doing that is to keep updating the values at those indices starting from the end of the array so that in each step after updating at some index i, the values of prefix sums for indices from 1 to i-1 won’t get affected and in each step we are actually decreasing the possible values of i for which A_1 +A_2 +...+ A_i>0 by 1.
Finally, we will have exactly K values of i where A_1+A_2+... + A_i>0 .
TIME COMPLEXITY:
O(n) for each testcase.
SOLUTION:
Editorialist's solution
#include <bits/stdc++.h>
using namespace std;
int main()
{
int t;
cin >> t;
while (t--)
{
int n, k;
cin >> n >> k;
vector<int> a(n);
for (int i = 0; i < n; i++)
{
if (i % 2 == 0)
a[i] = i + 1;
else
a[i] = -i - 1;
}
int pos = (n / 2) + (n % 2);
int neg = n / 2;
if (pos > k)
{
for (int i = n - 1; i >= 0; i--)
{
if (pos == k)
break;
if (a[i] > 0)
{
a[i] = -a[i];
pos--;
}
}
}
else
{
for (int i = n - 1; i >= 0; i--)
{
if (pos == k)
break;
if (a[i] < 0)
{
a[i] = -a[i];
pos++;
}
}
}
for (int i = 0; i < n; i++)
cout << a[i] << " ";
cout << endl;
}
return 0;
}
Setter's solution
#include <cmath>
#include <functional>
#include <fstream>
#include <iostream>
#include <vector>
#include <algorithm>
#include <string>
#include <set>
#include <map>
#include <list>
#include <time.h>
#include <math.h>
#include <random>
#include <deque>
#include <queue>
#include <cassert>
#include <unordered_map>
#include <unordered_set>
#include <iomanip>
#include <bitset>
#include <sstream>
#include <chrono>
#include <cstring>
using namespace std;
typedef long long ll;
#ifdef iq
mt19937 rnd(228);
#else
mt19937 rnd(chrono::high_resolution_clock::now().time_since_epoch().count());
#endif
int main()
{
#ifdef iq
freopen("a.in", "r", stdin);
#endif
ios::sync_with_stdio(0);
cin.tie(0);
int t;
cin >> t;
while (t--)
{
int n, k;
cin >> n >> k;
vector<int> a(n);
for (int i = 0; i < n; i++)
{
a[i] = i + 1;
if (i % 2 == 0)
{
a[i] *= -1;
}
}
int s = 0;
int cnt = 0;
vector<int> p(n);
for (int i = 0; i < n; i++)
{
s += a[i];
p[i] = s;
if (p[i] > 0)
cnt++;
}
bool gg = false;
auto print = [&](vector<int> &p) {
if (gg)
return;
for (int x : p)
cout << x << ' ';
cout << '\n';
gg = true;
};
if (cnt == k)
{
print(a);
}
int need = (cnt < k);
for (int i = n - 1; i >= 0; i--)
{
if (a[i] > 0 && need == 0)
{
for (int j = i; j < n; j++)
{
if (p[j] > 0)
cnt--;
p[j] -= 2 * a[i];
if (p[j] > 0)
cnt++;
}
a[i] *= -1;
if (cnt == k)
{
print(a);
break;
}
}
else if (a[i] < 0 && need == 1)
{
for (int j = i; j < n; j++)
{
if (p[j] > 0)
cnt--;
p[j] -= 2 * a[i];
if (p[j] > 0)
cnt++;
}
a[i] *= -1;
if (cnt == k)
{
print(a);
break;
}
}
}
}
}
VIDEO EDITORIAL:
Please comment below if you have any questions, alternate solutions, or suggestions.