Submission #14502304


Source Code Expand

Copy
import sys
import numpy as np

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

MOD = 998_244_353

def precompute(S):
    if np.all(S == 1):
        print(1)
        exit()
    while S[0] == 1:
        S = S[1:]
    if np.all(S == 0):
        print(1)
        exit()
    # 各 1 の左側にある 0 の個数を数える
    z = 0
    A = []
    for i, x in enumerate(S):
        if x == 0:
            z += 1
        else:
            A.append(z)
    return np.array(A, dtype=np.int64)

S, K = read().split()
S = np.array(list(S), np.int64) - ord('0')
K = int(K)
A = precompute(S)

A = precompute(S)

def main(A):
    U = 310
    # 数、消費回数、できた禁止地点
    dp = np.zeros((U, U, U), np.int64)
    dp[0, 0, 0] = 1
    for n in range(1, U):
        # 消費回数なし
        dp[n] += dp[n - 1]
        count = np.sum(A == n)
        if not count:
            continue
        tmp = np.zeros_like(dp[n])
        tmp[1:, 1:] += dp[n - 1][:-1, :-1]
        np.cumsum(tmp, axis=1, out=tmp)
        tmp %= MOD
        for k in range(1, count + 1):
            dp[n] += tmp
            tmp2 = np.zeros_like(tmp)
            tmp2[1:] = tmp[:-1]
            np.cumsum(tmp2, axis=1, out=tmp2)
            tmp = tmp2 % MOD
        dp[n, :, n + 1:] = 0
        dp[n] %= MOD
    x = dp[-1][:K + 1, :].sum() % MOD
    return x

print(main(A))

Submission Info

Submission Time
Task C - Shift
User maspy
Language Python (3.8.2)
Score 800
Code Size 1471 Byte
Status
Exec Time 848 ms
Memory 261244 KB

Judge Result

Set Name Score / Max Score Test Cases
Sample 0 / 0 s1.txt, s2.txt, s3.txt
All 800 / 800 01.txt, 02.txt, 03.txt, 04.txt, 05.txt, 06.txt, 07.txt, 08.txt, 09.txt, 10.txt, 11.txt, 12.txt, 13.txt, 14.txt, 15.txt, 16.txt, 17.txt, 18.txt, 19.txt, 20.txt, 21.txt, 22.txt, 23.txt, 24.txt, 25.txt, 26.txt, 27.txt, 28.txt, 29.txt, 30.txt, 31.txt, 32.txt, 33.txt, 34.txt, 35.txt, 36.txt, s1.txt, s2.txt, s3.txt
Case Name Status Exec Time Memory
01.txt 668 ms 260956 KB
02.txt 690 ms 260960 KB
03.txt 679 ms 261124 KB
04.txt 671 ms 261044 KB
05.txt 660 ms 261108 KB
06.txt 669 ms 260804 KB
07.txt 646 ms 260996 KB
08.txt 661 ms 261156 KB
09.txt 128 ms 27012 KB
10.txt 124 ms 27232 KB
11.txt 848 ms 260712 KB
12.txt 836 ms 260980 KB
13.txt 550 ms 261116 KB
14.txt 530 ms 261052 KB
15.txt 580 ms 261244 KB
16.txt 588 ms 261044 KB
17.txt 337 ms 260640 KB
18.txt 607 ms 261020 KB
19.txt 673 ms 260964 KB
20.txt 657 ms 260940 KB
21.txt 531 ms 260928 KB
22.txt 674 ms 260900 KB
23.txt 410 ms 261240 KB
24.txt 415 ms 261144 KB
25.txt 512 ms 260944 KB
26.txt 503 ms 261192 KB
27.txt 501 ms 261160 KB
28.txt 119 ms 27100 KB
29.txt 125 ms 26960 KB
30.txt 118 ms 27096 KB
31.txt 122 ms 26948 KB
32.txt 120 ms 27356 KB
33.txt 355 ms 261176 KB
34.txt 509 ms 260948 KB
35.txt 632 ms 260944 KB
36.txt 326 ms 260324 KB
s1.txt 341 ms 260940 KB
s2.txt 344 ms 261188 KB
s3.txt 426 ms 260880 KB