#942048 (Nim) No.2613 Sum of Combination

提出ソース

結果

問題	No.2613 Sum of Combination
ユーザー	Mitarushi
提出日時	2024-01-06 13:17:04
言語	Nim (2.0.2)
結果	AC
実行時間	1,248 ms / 4,500 ms
コード長	3,614 bytes
コンパイル時間	3,867 ms
コンパイル使用メモリ	66,096 KB
実行使用メモリ	14,416 KB
最終ジャッジ日時	2024-01-19 20:50:48
合計ジャッジ時間	33,203 ms
ジャッジサーバーID （参考情報）	judge12 / judge13

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 6,676 KB
testcase_01	AC	2 ms 6,676 KB
testcase_02	AC	19 ms 6,676 KB
testcase_03	AC	2 ms 6,676 KB
testcase_04	AC	2 ms 6,676 KB
testcase_05	AC	2 ms 6,676 KB
testcase_06	AC	2 ms 6,676 KB
testcase_07	AC	2 ms 6,676 KB
testcase_08	AC	3 ms 6,676 KB
testcase_09	AC	2 ms 6,676 KB
testcase_10	AC	2 ms 6,676 KB
testcase_11	AC	2 ms 6,676 KB
testcase_12	AC	2 ms 6,676 KB
testcase_13	AC	31 ms 6,676 KB
testcase_14	AC	31 ms 6,676 KB
testcase_15	AC	19 ms 6,676 KB
testcase_16	AC	31 ms 6,676 KB
testcase_17	AC	31 ms 6,676 KB
testcase_18	AC	31 ms 6,676 KB
testcase_19	AC	32 ms 6,676 KB
testcase_20	AC	3 ms 6,676 KB
testcase_21	AC	3 ms 6,676 KB
testcase_22	AC	65 ms 6,676 KB
testcase_23	AC	1,245 ms 13,872 KB
testcase_24	AC	1,239 ms 13,856 KB
testcase_25	AC	1,236 ms 13,264 KB
testcase_26	AC	1,233 ms 14,264 KB
testcase_27	AC	563 ms 8,704 KB
testcase_28	AC	1,225 ms 14,392 KB
testcase_29	AC	1,207 ms 14,008 KB
testcase_30	AC	1,230 ms 14,272 KB
testcase_31	AC	1,231 ms 14,032 KB
testcase_32	AC	1,224 ms 13,896 KB
testcase_33	AC	1,225 ms 14,304 KB
testcase_34	AC	1,207 ms 14,304 KB
testcase_35	AC	1,236 ms 14,304 KB
testcase_36	AC	1,235 ms 14,304 KB
testcase_37	AC	1,239 ms 14,304 KB
testcase_38	AC	1,227 ms 14,360 KB
testcase_39	AC	1,205 ms 14,384 KB
testcase_40	AC	1,220 ms 14,376 KB
testcase_41	AC	1,222 ms 14,416 KB
testcase_42	AC	1,248 ms 14,408 KB
testcase_43	AC	1,204 ms 14,304 KB
testcase_44	AC	1,047 ms 14,304 KB
testcase_45	AC	1 ms 6,676 KB
testcase_46	AC	2 ms 6,676 KB
testcase_47	AC	2 ms 6,676 KB
testcase_48	AC	2 ms 6,676 KB
testcase_49	AC	2 ms 6,676 KB
testcase_50	AC	694 ms 14,304 KB
testcase_51	AC	680 ms 14,304 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

import strutils, sequtils

iterator readInt(): int {.closure.} =
    while true:
        let line = stdin.readLine.split.map(parseInt)
        for x in line:
            yield x

const MOD = 998244353
func powMod(a, b, m: int): int =
    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: int): int =
    var phi = n - 1
    var factors = newSeq[int]()
    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[int]) =
    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[int]) =
    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() =
    var input = readInt
    var n = input()
    let p = input()

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

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

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

        var a = newSeq[int](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()

yukicoder

結果

テストケース

ソースコード