Yet Another Inversions Problem (Codeforces 1917D)
Tags: Math, number theory

You are given a permutation $$$p_0, p_1, \ldots, p_{n-1}$$$ of odd integers from $$$1$$$ to $$$2n-1$$$ and a permutation $$$q_0, q_1, \ldots, q_{k-1}$$$ of integers from $$$0$$$ to $$$k-1$$$.An array $$$a_0, a_1, \ldots, a_{nk-1}$$$ of length $$$nk$$$ is defined as follows: $$$a_{i \cdot k+j}=p_i \cdot 2^{q_j}$$$ for all $$$0 \le i < n$$$ and all $$$0 \le j < k$$$ For example, if $$$p = [3, 5, 1]$$$ and $$$q = [0, 1]$$$, then $$$a = [3, 6, 5, 10, 1, 2]$$$.Note that all arrays in the statement are zero-indexed. Note that each element of the array $$$a$$$ is uniquely determined.Find the number of inversions in the array $$$a$$$. Since this number can be very large, you should find only its remainder modulo $$$998\,244\,353$$$.An inversion in array $$$a$$$ is a pair $$$(i, j)$$$ ($$$0 \le i < j < nk$$$) such that $$$a_i > a_j$$$.

Input: The first line contains a single integer $$$t$$$ ($$$1 \le t \le 10^4$$$) — the number of test cases.The first line of each test case contains two integers $$$n$$$ and $$$k$$$ ($$$1 \le n, k \le 2 \cdot 10^5$$$) — the lengths of arrays $$$p$$$ and $$$q$$$.The second line of each test case contains $$$n$$$ distinct integers $$$p_0, p_1, \ldots, p_{n-1}$$$ ($$$1 \le p_i \le 2n-1$$$, $$$p_i$$$ is odd) — the array $$$p$$$.The third line of each test case contains $$$k$$$ distinct integers $$$q_0, q_1, \ldots, q_{k-1}$$$ ($$$0 \le q_i < k$$$) — the array $$$q$$$.It is guaranteed that the sum of $$$n$$$ over all test cases doesn't exceed $$$2 \cdot 10^5$$$ and the sum of $$$k$$$ over all test cases doesn't exceed $$$2 \cdot 10^5$$$.

Output: For each test case, output one integer: the number of inversions in array $$$a$$$ modulo $$$998\,244\,353$$$.

Note: In the first test case, array $$$a$$$ is equal to $$$[3, 6, 5, 10, 1, 2]$$$. There are $$$9$$$ inversions in it: $$$(0, 4)$$$, $$$(0, 5)$$$, $$$(1, 2)$$$, $$$(1, 4)$$$, $$$(1, 5)$$$, $$$(2, 4)$$$, $$$(2, 5)$$$, $$$(3, 4)$$$, $$$(3, 5)$$$. Note that these are pairs $$$(i, j)$$$ such that $$$i < j$$$ and $$$a_i > a_j$$$.In the second test case, array $$$a$$$ is equal to $$$[8, 4, 1, 2, 24, 12, 3, 6, 40, 20, 5, 10]$$$. There are $$$25$$$ inversions in it.In the third test case, array $$$a$$$ is equal to $$$[1, 2, 4, 8, 16]$$$. There are no inversions in it.

Sample Inputs:
4
3 2
3 5 1
0 1
3 4
1 3 5
3 2 0 1
1 5
1
0 1 2 3 4
8 3
5 1 7 11 15 3 9 13
2 0 1

Sample Outputs:
9
25
0
104

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

def cal(k, diff):
    if diff == 0:
        return k * (k + 1) // 2
    elif diff > 0:
        newk = max(k - diff, 0)
        return newk * (newk + 1) // 2
    else:
        actualdiff = -diff
        if actualdiff >= k - 1:
            return k * k
        temp = actualdiff * k + (1 + actualdiff + k) * (k - actualdiff) // 2
        return temp

def lso(index):
    return index - (index & index - 1)

def query(tree, index):
    s = 0
    l = len(tree)
    if index >= l:
        index = l - 1
    while index > 0:
        s += tree[index]
        index -= lso(index)
    return s

def update(tree, index, value):
    l = len(tree)
    while index < l:
        tree[index] += value
        index += lso(index)

def getUp(num, diff):
    if diff <= 0:
        return num * 2 ** abs(diff)
    else:
        actualdiff = diff - 1
        return num // 2 ** actualdiff

def getBottom(num, diff):
    if diff == 0:
        return num
    if diff < 0:
        return num * 2 ** (abs(diff) - 1)
    actualdiff = diff
    return num // 2 ** actualdiff
t = int(input())
for i in range(t):
    (n, k) = map(int, input().split())
    p = list(map(int, input().split()))
    q = list(map(int, input().split()))
    tree = [0 for j in range(k + 1)]
    ans = 0
    for num in q:
        ans += query(tree, num + 1)
        update(tree, num + 1, 1)
    ans = k * (k - 1) // 2 - ans
    ans *= n
    tree = [0 for j in range(2 * n + 1)]
    for num in p:
        diff = 0
        bottom = num
        up = num * 2
        while True:
            count = query(tree, up) - query(tree, bottom)
            ans += count * cal(k, diff)
            ans %= 998244353
            bottom *= 2
            up *= 2
            diff -= 1
            if bottom > 2 * n - 1:
                break
        diff = 1
        while True:
            up = getUp(num, diff)
            bottom = getBottom(num, diff)
            count = query(tree, up) - query(tree, bottom)
            ans += count * cal(k, diff)
            ans %= 998244353
            diff += 1
            if up < 1:
                break
        update(tree, num, 1)
    print(ans % 998244353)

Solution (Python):
from sys import *
input = stdin.readline
inputlist = lambda : map(int, stdin.readline().split())
from math import *
MOD = 998244353

class FenwickTree:

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

    @classmethod
    def from_array(cls, arr):
        tree = FenwickTree(len(arr))
        for (i, v) in enumerate(arr):
            tree.update(i, v)
        return tree

    def update(self, idx, delta):
        """"""
        while idx < self.n:
            self.bit[idx] += delta
            idx |= idx + 1

    def sum(self, right_idx):
        """"""
        res = 0
        while right_idx >= 0:
            res += self.bit[right_idx]
            right_idx = (right_idx & right_idx + 1) - 1
        return res

    def sum2(self, left_index, right_idx):
        """"""
        return self.sum(right_idx) - self.sum(left_index - 1)

def solve():
    (n, k) = inputlist()
    p = list(inputlist())
    q = list(inputlist())
    answer = 0
    tree = FenwickTree(2 * n)
    for i in range(1, 2 * n, 2):
        tree.update(i, 1)
    for i in p:
        (left, right) = (i + 1, min(2 * n - 1, 2 * i))
        diff = 1
        while left != 2 * n - 1 and diff < k:
            count = tree.sum2(left, right)
            answer = (answer + count * (k - diff) * (k - diff + 1) // 2) % MOD
            (left, right) = (right, min(2 * n - 1, 2 * right))
            diff += 1
        tree.update(i, -1)
    for i in range(1, 2 * n, 2):
        tree.update(i, 1)
    for i in reversed(p):
        (left, right) = (i + 1, min(2 * n - 1, 2 * i))
        diff = 1
        while left != 2 * n - 1:
            count = tree.sum2(left, right)
            if diff < k:
                answer = (answer + count * (k * k - (k - diff) * (k - diff + 1) // 2)) % MOD
            else:
                answer = answer + count * k * k % MOD
            (left, right) = (right, min(2 * n - 1, 2 * right))
            diff += 1
        tree.update(i, -1)
    temp = 0
    tree = FenwickTree(k)
    for i in q:
        temp += tree.sum2(i, k - 1)
        tree.update(i, 1)
    return (answer + temp * n) % MOD
for _ in range(int(input())):
    print(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()

class BIT:

    def __init__(self, n):
        self.bit = [0] * (n + 1)
        self.n = n

    def update(self, idx, val):
        idx += 1
        while idx <= self.n:
            self.bit[idx] += val
            idx += idx & -idx

    def prefixSum(self, idx):
        idx += 1
        ans = 0
        while idx > 0:
            ans += self.bit[idx]
            idx -= idx & -idx
        return ans

    def rangeSum(self, l, r):
        if l > r:
            return 0
        else:
            return self.prefixSum(r) - self.prefixSum(l - 1)

def solve():
    (n, k) = readInts()
    (p, q) = (readList(), readList())
    (ftp, ftq) = (BIT(2 * n + 2), BIT(k))
    ans = 0
    for i in range(k):
        ans += i - ftq.prefixSum(q[i])
        ftq.update(q[i], 1)
    ans *= n
    ans %= m
    for i in range(n):
        if k > 20:
            ans += i * ((k - 21) * k - (k + 20) * (k - 21) // 2)
            ans %= m
        ans += (i - ftp.prefixSum(p[i])) * k
        ans %= m
        for j in range(1, min(21, k)):
            ans += (i - ftp.prefixSum(p[i] // (1 << j)) + ftp.rangeSum(p[i] * (1 << j) + 1, 2 * n + 1)) * (k - j)
            ans %= m
        ftp.update(p[i], 1)
    return ans
m = 998244353
for _ in range(int(input())):
    print(solve())