Contest Duration: ~ (local time) (150 minutes) Back to Home

Submission #16917670

Source Code Expand

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

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)
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_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
# 追加処理
# 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

print(main(K, A))```

#### Submission Info

Submission Time 2020-09-20 22:26:21+0900 E - Shuffle Window maspy Python (3.8.2) 900 2499 Byte AC 685 ms 119476 KB

#### Judge Result

Set Name Sample All
Score / Max Score 0 / 0 900 / 900
Status
 AC × 2
 AC × 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