AtCoder Regular Contest 153

C - ± Increasing Sequence

### Problem Statement

You are given a sequence of length $N$, $A = (A_1, \ldots, A_N)$, consisting of $1$ and $-1$.

Determine whether there is an integer sequence $x = (x_1, \ldots, x_N)$ that satisfies all of the following conditions, and print one such sequence if it exists.

• $|x_i| \leq 10^{12}$ for every $i$ ($1\leq i\leq N$).
• $x$ is strictly increasing. That is, $x_1 < \cdots < x_N$.
• $\sum_{i=1}^N A_ix_i = 0$.

### Constraints

• $1\leq N\leq 2\times 10^5$
• $A_i \in \lbrace 1, -1\rbrace$

### Input

The input is given from Standard Input in the following format:

$N$
$A_1$ $\ldots$ $A_N$

### Output

If there is an integer sequence $x$ that satisfies all of the conditions in question, print Yes; otherwise, print No. In case of Yes, print the elements of such an integer sequence $x$ in the subsequent line, separated by spaces.

$x_1$ $\ldots$ $x_N$

If multiple integer sequences satisfy the conditions, you may print any of them.

5
-1 1 -1 -1 1

### Sample Output 1

Yes
-3 -1 4 5 7

For this output, we have $\sum_{i=1}^NA_ix_i= -(-3) + (-1) - 4 - 5 + 7 = 0$.

1
-1

Yes
0

2
1 -1

No

### 题解

• 直接用 $x_i=i$ 作为初始解，算出 $sum=\sum_{i=1}^N A_ix_i$
• 如果 $sum>0$ 那么想办法让 $sum$ 减小为 $0$
• 如果 $sum<0$ 那么想办法让 $sum$ 增大为 $0$

• 将第 $k,k+1,\dots,N$ 位同时增大 $t$（$k,t\in\mathbb{Z}\ 且\ 1\leq k\leq N$）
• 将第 $1,2,\dots,k$ 位同时减小 $t$（$k,t\in\mathbb{Z}\ 且\ 1\leq k\leq N$）

$$1,2,\dots,k-1,k+t,k+1+t,\dots,N-1+t,N+t(第一种操作)\\ 或者\\ 1-t,2-t,\dots,k-1-t,k-t,k+1,\dots,N-1,N(第二种操作)$$

### 代码

#include <bits/stdc++.h>
#define int long long
#define endl '\n'

using namespace std;

signed main()
{
ios_base::sync_with_stdio(false);
cin.tie(0);
cout.tie(0);
int N;
cin >> N;
vector<int> A(N + 1);
for (int i = 1; i <= N; i++)
cin >> A[i];
int sum = 0;
for (int i = 1; i <= N; i++)
sum += A[i] * i;
if (sum > 0)
{
int cnt_posi = 0, cnt_nega = 0, pos = 0;
for (int i = N; i >= 1; i--)
{
if (A[i] == 1)
cnt_posi++;
else
cnt_nega++;
if (cnt_nega > cnt_posi)
{
pos = i;
break;
}
}
if (pos)
{
cout << "Yes" << endl;
for (int i = 1; i < pos; i++)
cout << i << ' ';
for (int i = pos; i <= N; i++)
cout << i + sum << ' ';
cout << endl;
return 0;
}
cnt_posi = 0, cnt_nega = 0;
for (int i = 1; i <= N; i++)
{
if (A[i] == 1)
cnt_posi++;
else
cnt_nega++;
if (cnt_posi > cnt_nega)
{
pos = i;
break;
}
}
if (pos)
{
cout << "Yes" << endl;
for (int i = 1; i <= pos; i++)
cout << i - sum << ' ';
for (int i = pos + 1; i <= N; i++)
cout << i << ' ';
cout << endl;
return 0;
}
}
else
{
int cnt_posi = 0, cnt_nega = 0, pos = 0;
for (int i = N; i >= 1; i--)
{
if (A[i] == 1)
cnt_posi++;
else
cnt_nega++;
if (cnt_nega < cnt_posi)
{
pos = i;
break;
}
}
if (pos)
{
cout << "Yes" << endl;
for (int i = 1; i < pos; i++)
cout << i << ' ';
for (int i = pos; i <= N; i++)
cout << i - sum << ' ';
cout << endl;
return 0;
}
cnt_posi = 0, cnt_nega = 0;
for (int i = 1; i <= N; i++)
{
if (A[i] == 1)
cnt_posi++;
else
cnt_nega++;
if (cnt_posi < cnt_nega)
{
pos = i;
break;
}
}
if (pos)
{
cout << "Yes" << endl;
for (int i = 1; i <= pos; i++)
cout << i + sum << ' ';
for (int i = pos + 1; i <= N; i++)
cout << i << ' ';
cout << endl;
return 0;
}
}
cout << "No" << endl;
return 0;
}