結果

問題 No.2613 Sum of Combination
ユーザー MitarushiMitarushi
提出日時 2024-01-06 03:09:24
言語 Nim
(2.0.2)
結果
WA  
実行時間 -
コード長 3,517 bytes
コンパイル時間 3,349 ms
コンパイル使用メモリ 65,968 KB
実行使用メモリ 14,412 KB
最終ジャッジ日時 2024-01-19 20:50:41
合計ジャッジ時間 28,142 ms
ジャッジサーバーID
（参考情報） 		judge11 / judge15
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 1 ms
6,676 KB
testcase_01 AC 2 ms
6,676 KB
testcase_02 WA -
testcase_03 AC 1 ms
6,676 KB
testcase_04 AC 1 ms
6,676 KB
testcase_05 AC 2 ms
6,676 KB
testcase_06 AC 1 ms
6,676 KB
testcase_07 AC 1 ms
6,676 KB
testcase_08 WA -
testcase_09 WA -
testcase_10 AC 2 ms
6,676 KB
testcase_11 WA -
testcase_12 WA -
testcase_13 WA -
testcase_14 WA -
testcase_15 WA -
testcase_16 WA -
testcase_17 WA -
testcase_18 WA -
testcase_19 WA -
testcase_20 WA -
testcase_21 WA -
testcase_22 WA -
testcase_23 WA -
testcase_24 WA -
testcase_25 WA -
testcase_26 WA -
testcase_27 WA -
testcase_28 WA -
testcase_29 WA -
testcase_30 WA -
testcase_31 WA -
testcase_32 WA -
testcase_33 WA -
testcase_34 WA -
testcase_35 WA -
testcase_36 WA -
testcase_37 WA -
testcase_38 WA -
testcase_39 WA -
testcase_40 WA -
testcase_41 WA -
testcase_42 WA -
testcase_43 WA -
testcase_44 WA -
testcase_45 AC 2 ms
6,676 KB
testcase_46 AC 2 ms
6,676 KB
testcase_47 WA -
testcase_48 WA -
testcase_49 WA -
testcase_50 AC 597 ms
14,300 KB
testcase_51 AC 602 ms
14,300 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

import strutils, sequtils

const MOD = 998244353
func powMod(a, b, m: int64): int64 =
    result = 1
    var 
        a = a mod m
        b = b
    while b > 0:
        if (b and 1) == 1:
            result = result * a mod m
        a = a * a mod m
        b = b shr 1

func findPrimitiveRoot(n: int64): int64 =
    var phi = n - 1
    var factors = newSeq[int64](0)
    for i in 2..n:
        if i * i > phi:
            break
        if phi mod i == 0:
            factors.add(i)
            while phi mod i == 0:
                phi = phi div i
    if phi > 1:
        factors.add(phi)
    for res in 1..<n:
        var ok = true
        for factor in factors:
            if powMod(res, (n - 1) div factor, n) == 1:
                ok = false
                break
        if ok:
            return res
    -1

const primitiveRoot = findPrimitiveRoot(MOD)

proc ntt(a: var seq[int64]) =
    let n = a.len
    var m = n
    while m > 1:
        let mh = m shr 1
        let wm = powMod(primitiveRoot, (MOD - 1) div m, MOD)
        var w = 1
        for i in 0..<mh:
            for j in countup(i, n - 1, m):
                let k = j + mh
                let
                    a0 = a[j]
                    a1 = a[k]
                a[j] = a0 + a1
                if a[j] >= MOD:
                    a[j] -= MOD
                a[k] = (a0 - a1 + MOD) * w mod MOD
            w = w * wm mod MOD
        m = mh
    
proc intt(a: var seq[int64]) =
    let n = a.len
    var m = 2
    while m <= n:
        let mh = m shr 1
        let wm = powMod(primitiveRoot, MOD - 1 - (MOD - 1) div m, MOD)
        var w = 1
        for i in 0..<mh:
            for j in countup(i, n - 1, m):
                let k = j + mh
                let
                    a0 = a[j]
                    a1 = a[k] * w mod MOD
                a[j] = a0 + a1
                if a[j] >= MOD:
                    a[j] -= MOD
                a[k] = a0 - a1
                if a[k] < 0:
                    a[k] += MOD
            w = w * wm mod MOD
        m = m shl 1
    let inv = powMod(n, MOD - 2, MOD)
    for i in 0..<n:
        a[i] = a[i] * inv mod MOD

proc solve() =
    let input = stdin.readLine.split.map(parseBiggestInt)
    var n = input[0]
    let p = input[1]

    let q = findPrimitiveRoot(p)
    var indexTable = newSeq[int64](p)
    indexTable[0] = 0
    var k = 1
    for i in 0..<p - 1:
        indexTable[k] = i
        k = k * q mod p
    var indexTableSum = newSeq[int64](p)
    for i in 1..<p:
        indexTableSum[i] = indexTableSum[i - 1] + indexTable[i]
        if indexTableSum[i] >= p:
            indexTableSum[i] -= p
    
    func comb(a, b: int64): int64 =
        (indexTableSum[a] - indexTableSum[b] - indexTableSum[a - b] + 3 * (p - 1)) mod (p - 1)

    var len = 1
    while len < (p - 1) * 2:
        len *= 2

    var count = newSeq[int64](len)
    count[0] = 1
    while n > 0:
        let m = n mod p
        n = n div p

        var a = newSeq[int64](len)
        for i in 0..m:
            a[comb(m, i)] += 1
        
        ntt(count)
        ntt(a)

        for i in 0..<len:
            count[i] = count[i] * a[i] mod MOD
        
        intt(count)

        for i in countdown(len - 1, p - 1):
            count[i - p + 1] += count[i]
            if count[i - p + 1] >= MOD:
                count[i - p + 1] -= MOD
            count[i] = 0
    
    var ans = 0
    k = 1
    for i in 0..<p - 1:
        ans = (ans + count[i] * k) mod MOD
        k = k * q mod p
    
    echo ans

solve()
0