Imbalanced Arrays (Codeforces 1853D)
Tags: Greedy, two pointers

Ntarsis has come up with an array $$$a$$$ of $$$n$$$ non-negative integers.Call an array $$$b$$$ of $$$n$$$ integers imbalanced if it satisfies the following:   $$$-n\le b_i\le n$$$, $$$b_i \ne 0$$$,  there are no two indices $$$(i, j)$$$ ($$$1 \le i, j \le n$$$) such that $$$b_i + b_j = 0$$$,  for each $$$1 \leq i \leq n$$$, there are exactly $$$a_i$$$ indices $$$j$$$ ($$$1 \le j \le n$$$) such that $$$b_i+b_j>0$$$, where $$$i$$$ and $$$j$$$ are not necessarily distinct. Given the array $$$a$$$, Ntarsis wants you to construct some imbalanced array. Help him solve this task, or determine it is impossible.

Input: Each test contains multiple test cases. The first line contains the number of test cases $$$t$$$ ($$$1 \le t \le 10^5$$$). The description of the test cases follows.The first line of each test case has a single integer $$$n$$$ ($$$1 \leq n \leq 10^5$$$).The next line contains $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$ ($$$0 \leq a_i \leq n$$$).It is guaranteed that the sum of $$$n$$$ across all test cases does not exceed $$$10^5$$$.

Output: For each test case, output "NO" if there exists no imbalanced array. Otherwise, output "YES". Then, on the next line, output $$$n$$$ integers $$$b_1, b_2, \ldots, b_n$$$ where $$$b_i \neq 0$$$ for all $$$1 \leq i \leq n$$$ — an imbalanced array.

Note: For the first test case, $$$b = [1]$$$ is an imbalanced array. This is because for $$$i = 1$$$, there is exactly one $$$j$$$ ($$$j = 1$$$) where $$$b_1 + b_j > 0$$$.For the second test case, it can be shown that there exists no imbalanced array.For the third test case, $$$a = [0, 1, 0]$$$. The array $$$b = [-3, 1, -2]$$$ is an imbalanced array.   For $$$i = 1$$$ and $$$i = 3$$$, there exists no index $$$j$$$ such that $$$b_i + b_j > 0$$$.  For $$$i = 2$$$, there is only one index $$$j = 2$$$ such that $$$b_i + b_j > 0$$$ ($$$b_2 + b_2 = 1 + 1 = 2$$$).  Another possible output for the third test case could be $$$b = [-2, 1, -3]$$$.

Sample Inputs:
5
1
1
4
1 4 3 4
3
0 1 0
4
4 3 2 1
3
1 3 1

Sample Outputs:
YES
1 
NO
YES
-3 1 -2 
YES
4 2 -1 -3 
YES
-1 3 -1

Solution (Python):
import sys
input = sys.stdin.readline

def readList():
    return list(map(int, input().split()))

def readInt():
    return int(input())

def readInts():
    return map(int, input().split())

def readStr():
    return input().strip()

def solve():
    (n, arr) = (readInt(), readList())
    ans = [-1] * n
    arr = sorted([(v, i) for (i, v) in enumerate(arr)])
    (f, b, of) = (0, n - 1, 0)
    while f < b:
        (vb, vi) = arr[b]
        if vb - of == 0:
            break
        k = n - (vb - of)
        if k == 0:
            if arr[f][0] + of <= 0:
                print('NO')
                return
            ans[vi] = vb - of
            of += 1
            n -= 1
            b -= 1
        elif f + k - 1 >= b:
            print('NO')
            return
        elif not arr[f][0] == arr[f + k - 1][0] == of:
            print('NO')
            return
        elif arr[f + k][0] == of:
            print('NO')
            return
        else:
            for i in range(f, f + k):
                ans[arr[i][1]] = -(vb - of + 1)
            ans[vi] = vb - of
            of += 1
            n -= k + 1
            b -= 1
            f += k
    if f <= b:
        (v, i) = arr[f]
        v -= of
        if v < 0:
            print('NO')
            return
        for j in range(f, b + 1):
            ans[arr[j][1]] = v - 1 if v == 0 else v
    print('YES')
    print(*ans)
    return
for _ in range(int(input())):
    solve()

Solution (Python):
import sys
input = sys.stdin.readline

def readList():
    return list(map(int, input().split()))

def readInt():
    return int(input())

def readInts():
    return map(int, input().split())

def readStr():
    return input().strip()

def solve():
    (n, arr) = (readInt(), readList())
    ans = [-1] * n
    arr = sorted([(v, i) for (i, v) in enumerate(arr)])
    (f, b, of) = (0, n - 1, 0)
    while f < b:
        (vb, vi) = arr[b]
        vb -= of
        if vb == 0:
            break
        k = n - vb
        ans[vi] = vb
        if k == 0:
            if arr[f][0] + of <= 0:
                print('NO')
                return
        elif f + k - 1 < b and arr[f][0] == arr[f + k - 1][0] == of and (arr[f + k][0] > of):
            for i in range(f, f + k):
                ans[arr[i][1]] = -vb - 1
        else:
            print('NO')
            return
        of += 1
        n -= k + 1
        b -= 1
        f += k
    if f <= b:
        (v, i) = arr[f]
        v -= of
        if v < 0:
            print('NO')
            return
        for j in range(f, b + 1):
            ans[arr[j][1]] = v - 1 if v == 0 else v
    print('YES')
    print(*ans)
    return
for _ in range(int(input())):
    solve()

Solution (Python):
import sys
input = sys.stdin.readline

def readList():
    return list(map(int, input().split()))

def readInt():
    return int(input())

def readInts():
    return map(int, input().split())

def readStr():
    return input().strip()

def solve():
    (n, arr) = (readInt(), readList())
    ans = [-1] * n
    arr = sorted([(v, i) for (i, v) in enumerate(arr)])
    (f, b, of) = (0, n - 1, 0)
    while f <= b:
        (vb, vi) = arr[b]
        vb -= of
        if vb == 0:
            (v, i) = arr[f]
            v -= of
            if v < 0:
                print('NO')
                return
            for j in range(f, b + 1):
                ans[arr[j][1]] = v - 1 if v == 0 else v
            break
        k = n - vb
        ans[vi] = vb
        if k == 0:
            if arr[f][0] + of <= 0:
                print('NO')
                return
        elif f + k - 1 < b and arr[f][0] == arr[f + k - 1][0] == of and (arr[f + k][0] > of):
            for i in range(f, f + k):
                ans[arr[i][1]] = -vb - 1
        else:
            print('NO')
            return
        of += 1
        n -= k + 1
        b -= 1
        f += k
    print('YES')
    print(*ans)
    return
for _ in range(int(input())):
    solve()