Inversion Composition (Codeforces 1951F)
Tags: Greedy, data structures

My Chemical Romance - DisenchantedඞYou are given a permutation $$$p$$$ of size $$$n$$$, as well as a non-negative integer $$$k$$$. You need to construct a permutation $$$q$$$ of size $$$n$$$ such that $$$\operatorname{inv}(q) + \operatorname{inv}(q \cdot p) = k {}^\dagger {}^\ddagger$$$, or determine if it is impossible to do so.$$${}^\dagger$$$ For two permutations $$$p$$$ and $$$q$$$ of the same size $$$n$$$, the permutation $$$w = q \cdot p$$$ is such that $$$w_i = q_{p_i}$$$ for all $$$1 \le i \le n$$$.$$${}^\ddagger$$$ For a permutation $$$p$$$ of size $$$n$$$, the function $$$\operatorname{inv}(p)$$$ returns the number of inversions of $$$p$$$, i.e. the number of pairs of indices $$$1 \le i < j \le n$$$ such that $$$p_i > p_j$$$.

Input: Each test contains multiple test cases. The first line contains an integer $$$t$$$ ($$$1 \le t \le 10^4$$$) — the number of test cases. The description of the test cases follows.The first line of each test case contains two integers $$$n$$$ and $$$k$$$ ($$$1 \le n \le 3 \cdot 10^5, 0 \le k \le n(n - 1)$$$) — the size of $$$p$$$ and the target number of inversions.The second line of each test case contains $$$n$$$ integers $$$p_1, p_2, \ldots, p_n$$$ ($$$1 \le p_i \le n$$$, $$$p_i$$$'s are pairwise distinct) — the given permutation $$$p$$$.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$3 \cdot 10^5$$$.

Output: For each test case, print on one line "YES" if there exists a permutation $$$q$$$ that satisfies the given condition, or "NO" if there is no such permutation.If the answer is "YES", on the second line, print $$$n$$$ integers $$$q_1, q_2, \ldots, q_n$$$ that represent such a satisfactory permutation $$$q$$$. If there are multiple such $$$q$$$'s, print any of them.

Note: In the first test case, we have $$$q \cdot p = [2, 1, 3]$$$, $$$\operatorname{inv}(q) = 3$$$, and $$$\operatorname{inv}(q \cdot p) = 1$$$.In the fourth test case, we have $$$q \cdot p = [9, 1, 8, 5, 7, 6, 4, 3, 2]$$$, $$$\operatorname{inv}(q) = 24$$$, and $$$\operatorname{inv}(q \cdot p) = 27$$$.

Sample Inputs:
5
3 4
2 3 1
5 5
2 3 5 1 4
6 11
5 1 2 3 4 6
9 51
3 1 4 2 5 6 7 8 9
1 0
1

Sample Outputs:
YES
3 2 1
NO
NO
YES
1 5 9 8 7 6 4 3 2
YES
1

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

class FenwickTree:

    def __init__(self, n):
        self.size = n
        self.data = [0] * n

    def update(self, idx, x):
        while idx < self.size:
            self.data[idx] += x
            idx |= idx + 1

    def query(self, end):
        x = 0
        while end:
            x += self.data[end - 1]
            end &= end - 1
        return x

def solve():
    (n, k) = map(int, input().split())
    P = [int(x) - 1 for x in map(int, input().split())]
    ft = FenwickTree(n)
    pinv = n * (n - 1) // 2
    X = []
    for p in P:
        x = ft.query(p)
        pinv -= x
        X.append(x)
        ft.update(p, 1)
    if k < pinv or k > n * (n - 1) - pinv or (k - pinv) % 2 == 1:
        return 'NO'
    k = (k - pinv) // 2
    r = n
    for (i, x) in enumerate(X):
        if k - x <= 0:
            r = i
            break
        k -= x
    W = list(range(r, -1, -1)) + list(range(r + 1, n))
    l = r
    for i in range(r):
        if P[i] < P[r]:
            k -= 1
        if k == 0:
            l = i
            break
    for i in range(l + 1, r):
        (W[i], W[i + 1]) = (W[i + 1], W[i])
    Q = [0] * n
    for (w, p) in zip(W, P):
        Q[p] = w
    ans = ' '.join([str(q + 1) for q in Q])
    return f'YES\n{ans}'
T = int(input())
out = [solve() for _ in range(T)]
print('\n'.join(map(str, out)))

Solution (Python):
import sys
input = sys.stdin.readline
F = lambda : [*map(int, input().split())]
for _ in range(int(input())):
    (N, K) = F()
    A = F()
    POS = [0] * N
    for i in range(N):
        A[i] -= 1
        POS[A[i]] = i
    CNT = [0] * N
    SEGT = [0] * (2 * N)
    for a in A:
        l = N
        r = a + N - 1
        cnt = 0
        while l <= r:
            if l & 1 == 1:
                cnt += SEGT[l]
            if r & 1 == 0:
                cnt += SEGT[r]
            l = l + 1 >> 1
            r = r - 1 >> 1
        CNT[a] = cnt
        p = a + N
        while p:
            SEGT[p] += 1
            p >>= 1
    tot = N * (N - 1) // 2
    s = sum(CNT)
    mn = tot - s
    mx = tot + s
    if K - mn & 1 or K < mn or K > mx:
        print('NO')
        continue
    ANS = [0] * N
    now = 1
    inv_cnt = mn
    i = N
    while inv_cnt < K:
        i -= 1
        add = CNT[i]
        if inv_cnt + add * 2 <= K:
            ANS[i] = now
            now += 1
            inv_cnt += 2 * add
            continue
        extra = add - (K - inv_cnt) // 2
        j = 0
        while extra:
            ANS[j] = now
            now += 1
            if POS[j] < POS[i]:
                extra -= 1
            j += 1
        ANS[i] = now
        now += 1
        break
    for i in range(N):
        if ANS[i] == 0:
            ANS[i] = now
            now += 1
    print('YES')
    print(*ANS)

Solution (Python):
import sys
input = sys.stdin.readline
F = lambda : [*map(int, input().split())]
for _ in range(int(input())):
    (N, K) = F()
    A = F()
    POS = [0] * N
    for i in range(N):
        A[i] -= 1
        POS[A[i]] = i
    CNT = [0] * N
    SEGT = [0] * (2 * N)
    for a in A:
        l = N
        r = a + N - 1
        cnt = 0
        while l <= r:
            if l & 1 == 1:
                cnt += SEGT[l]
            if r & 1 == 0:
                cnt += SEGT[r]
            l = l + 1 >> 1
            r = r - 1 >> 1
        CNT[a] = cnt
        p = a + N
        while p:
            SEGT[p] += 1
            p >>= 1
    tot = N * (N - 1) // 2
    s = sum(CNT)
    mn = tot - s
    mx = tot + s
    if K - mn & 1 or K < mn or K > mx:
        print('NO')
        continue
    ANS = [0] * N
    now = 1
    inv_cnt = mn
    i = N
    while inv_cnt < K:
        i -= 1
        add = CNT[i]
        if inv_cnt + add * 2 <= K:
            ANS[i] = now
            now += 1
            inv_cnt += 2 * add
            continue
        extra = add - (K - inv_cnt) // 2
        j = 0
        while extra:
            ANS[j] = now
            now += 1
            if POS[j] < POS[i]:
                extra -= 1
            j += 1
        ANS[i] = now
        now += 1
        break
    for i in range(N):
        if ANS[i] == 0:
            ANS[i] = now
            now += 1
    print('YES')
    print(*ANS)