Submission #16917670


Source Code Expand

Copy
import sys
import numpy as np
import numba
from numba import njit, b1, i4, i8, f8

read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

MOD = 998_244_353

@njit((i8, i8), cache=True)
def mpow(a, n):
    p = 1
    while n:
        if n & 1:
            p = p * a % MOD
        a = a * a % MOD
        n >>= 1
    return p

@njit((i8[:], ), cache=True)
def build(raw_data):
    bit = raw_data.copy()
    for i in range(len(bit)):
        j = i + (i & (-i))
        if j < len(bit):
            bit[j] += bit[i]
            bit[j] %= MOD
    return bit


@njit((i8[:], i8), cache=True)
def get_sum(bit, i):
    s = 0
    while i:
        s += bit[i]
        i -= i & -i
    return s % MOD


@njit((i8[:], i8, i8), cache=True)
def add(bit, i, x):
    while i < len(bit):
        bit[i] += x
        bit[i] %= MOD
        i += i & -i

@njit((i8, i8[:]), cache=True)
def main(K, A):
    N = len(A)
    # 各処理の段階で、n が active になっている確率を求める。
    # 個数の分布も持つ。
    # 更新を sparse にするために、時刻(右端のインデックス) i にアクティブになったところの
    # 生存率には、p^{-i} を入れておく。時刻 t には確率 p^{t-i} なので、和に p^t をかける
    p = (1 - mpow(K, MOD - 2)) % MOD
    p_inv = mpow(p, MOD - 2)
    bit_active = np.zeros(N + 1, np.int64)
    bit_cnt = np.zeros(N + 1, np.int64)
    half = (MOD + 1) // 2
    ans = 0
    # 最初の起動
    ans += K * (K - 1) // 2 % MOD * half % MOD
    ans %= MOD
    for i in range(K):
        x = A[i]
        add(bit_cnt, x, 1)
        add(bit_active, x, mpow(p_inv, K - 1))
    for i in range(K, N):
        x = A[i]
        low_cnt = get_sum(bit_cnt, x)
        high_cnt = i - low_cnt
        pt = mpow(p, i)
        active_sum_low = get_sum(bit_active, x) * pt % MOD
        active_sum_high = K - 1 - active_sum_low
        # 追加処理
        add(bit_cnt, x, 1)
        add(bit_active, x, mpow(p_inv, i))
        # A[i] と A[j] (j<i) の転倒数の期待値を計算
        inv = 0
        inv += active_sum_low * half % MOD
        inv += (high_cnt - active_sum_high)
        inv += active_sum_high * half % MOD
        inv %= MOD
        ans += inv
        ans %= MOD
    return ans

N, K = map(int, readline().split())
A = np.array(read().split(), np.int64)

print(main(K, A))

Submission Info

Submission Time
Task E - Shuffle Window
User maspy
Language Python (3.8.2)
Score 900
Code Size 2499 Byte
Status
Exec Time 685 ms
Memory 119476 KB

Judge Result

Set Name Sample All
Score / Max Score 0 / 0 900 / 900
Status
× 2
× 26
Set Name Test Cases
Sample example_00, example_01
All example_00, example_01, max_random_00, max_random_01, max_random_02, mid_k_00, mid_k_01, mid_k_02, random_00, random_01, random_02, small_00, small_01, small_02, small_03, small_04, small_05, small_06, small_07, small_08, small_09, small_k_00, small_k_01, small_k_02, small_k_03, small_k_04
Case Name Status Exec Time Memory
example_00 590 ms 106168 KB
example_01 562 ms 106188 KB
max_random_00 646 ms 119040 KB
max_random_01 655 ms 119476 KB
max_random_02 646 ms 118972 KB
mid_k_00 642 ms 118752 KB
mid_k_01 672 ms 118728 KB
mid_k_02 671 ms 119364 KB
random_00 606 ms 114772 KB
random_01 629 ms 116768 KB
random_02 574 ms 110028 KB
small_00 595 ms 106988 KB
small_01 558 ms 106120 KB
small_02 557 ms 106812 KB
small_03 571 ms 106448 KB
small_04 577 ms 105728 KB
small_05 562 ms 106200 KB
small_06 571 ms 106208 KB
small_07 575 ms 106180 KB
small_08 579 ms 106824 KB
small_09 548 ms 106276 KB
small_k_00 685 ms 118776 KB
small_k_01 649 ms 118732 KB
small_k_02 650 ms 118208 KB
small_k_03 656 ms 118252 KB
small_k_04 669 ms 118748 KB