結果

問題 No.2613 Sum of Combination
ユーザー Mitarushi
提出日時 2024-01-06 03:09:24
言語 Nim
(2.2.0)
結果
WA  
実行時間 -
コード長 3,517 bytes
コンパイル時間 3,269 ms
コンパイル使用メモリ 65,920 KB
実行使用メモリ 14,464 KB
最終ジャッジ日時 2024-09-28 03:46:43
合計ジャッジ時間 27,027 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2 WA * 1
other AC * 10 WA * 39
権限があれば一括ダウンロードができます

ソースコード

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