結果

問題 No.1307 Rotate and Accumulate
ユーザー Shinya Fujita
提出日時 2025-01-19 01:10:59
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 459 ms / 5,000 ms
コード長 5,562 bytes
コンパイル時間 401 ms
コンパイル使用メモリ 82,048 KB
実行使用メモリ 122,700 KB
最終ジャッジ日時 2025-01-19 01:11:08
合計ジャッジ時間 6,608 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 19
権限があれば一括ダウンロードができます

ソースコード

diff # 

N, Q = map(int, input().split())
A = list(map(int, input().split()))
R = list(map(int, input().split()))


def primitive_root(m):
    """
    整数mの最小原始根を計算する
    Parameters
    ----------
    m : int
        2以上の整数
    Returns
    -------
    int
        mの最小原始根
    """
    if m == 2:
        return 1
    if m == 167772161:
        return 3
    if m == 469762049:
        return 3
    if m == 754974721:
        return 11
    if m == 998244353:
        return 3
    if m == 1000000007:
        return 5

    divs = [2] + [0] * 19
    cnt = 1
    x = (m - 1) // 2
    while x % 2 == 0:
        x //= 2

    i = 3
    while i**2 <= x:
        if x % i == 0:
            divs[cnt] = i
            cnt += 1
            while x % i == 0:
                x //= i

        i += 2

    if x > 1:
        divs[cnt] = x
        cnt += 1

    g = 2
    while True:
        for i in range(cnt):
            if pow(g, (m-1)//divs[i], m) == 1:
                break

        else:
            return g

        g += 1


def ceil_pow2(n):
    """
    n <= 2**x を満たす最小のxを返却する。
    Parameters
    ----------
    n : int
        0以上の整数
    Returns
    -------
    int
        n <= 2**x を満たす最小のx
    """
    if n < 1:
        return 0
    return (n-1).bit_length()


def bsf(n):
    """
    自然数を2bitで表現したときに、右から見て最初に1が立つ桁が何桁目かを返却する。
    （0-indexed）
    Parameters
    ----------
    n : int
        1以上の整数
    Returns
    -------
    int
        2bitで表現したときに、右から見て最初に1が立つ桁（0-indexed）
    """
    return (n & -n).bit_length()-1


MOD = 998244353
class Convolution:
    """
    畳み込みを行う。
    長さNの数列Aと長さMの数列Bから、長さ(N+M-1)の数列Cを計算する。
    c_i = sum_{j=0}^{i} a_j * b_{i-j}
    """
    def __init__(self):
        self._first1 = True
        self._first2 = True
        self._sum_e = [0] * 30
        self._sum_ie = [0] * 30
        self._root = primitive_root(MOD)

    def _butterfly(self, a):
        n = len(a)
        h = ceil_pow2(n)
        if self._first1:
            self._first1 = False
            es = [0] * 30
            ies = [0] * 30
            m = MOD - 1
            cnt2 = bsf(m)
            e = pow(self._root, m >> cnt2, MOD)
            ie = pow(e, MOD - 2, MOD)
            for i in range(cnt2 - 1)[::-1]:
                es[i] = e
                ies[i] = ie
                e *= e
                e %= MOD
                ie *= ie
                ie %= MOD

            now = 1
            for i in range(cnt2 - 1):
                self._sum_e[i] = es[i] * now % MOD
                now *= ies[i]
                now %= MOD

        for ph in range(1, h + 1):
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            now = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    left = a[i + offset]
                    right = a[i + offset + p] * now
                    a[i + offset] = (left + right) % MOD
                    a[i + offset + p] = (left - right) % MOD

                now *= self._sum_e[bsf(~s)]
                now %= MOD

    def _butterfly_inv(self, a):
        n = len(a)
        h = ceil_pow2(n)
        if self._first2:
            self._first2 = False
            es = [0] * 30
            ies = [0] * 30
            m = MOD - 1
            cnt2 = bsf(m)
            e = pow(self._root, m >> cnt2, MOD)
            ie = pow(e, MOD - 2, MOD)
            for i in range(cnt2 - 1)[::-1]:
                es[i] = e
                ies[i] = ie
                e *= e
                e %= MOD
                ie *= ie
                ie %= MOD

            now = 1
            for i in range(cnt2 - 1):
                self._sum_ie[i] = ies[i] * now % MOD
                now *= es[i]
                now %= MOD

        for ph in range(1, h + 1)[::-1]:
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            inow = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    left = a[i + offset]
                    right = a[i + offset + p]
                    a[i + offset] = (left + right) % MOD
                    a[i + offset + p] = (
                        (MOD + left - right) * inow % MOD
                    )
                inow *= self._sum_ie[bsf(~s)]
                inow %= MOD

    def convolution(self, a, b):
        n = len(a)
        m = len(b)
        if n*m == 0:
            return []

        if min(n, m) <= 60:
            if n < m:
                n, m = m, n
                a, b = b, a

            res = [0] * (n + m - 1)
            for i in range(n):
                for j in range(m):
                    res[i + j] += a[i] * b[j]
                    res[i + j] %= MOD

            return res

        z = 1 << ceil_pow2(n + m - 1)
        a += [0] * (z - n)
        b += [0] * (z - m)
        self._butterfly(a)
        self._butterfly(b)
        for i in range(z):
            a[i] *= b[i]
            a[i] %= MOD

        self._butterfly_inv(a)
        a = a[:n + m - 1]
        iz = pow(z, MOD - 2, MOD)
        for i in range(n + m - 1):
            a[i] *= iz
            a[i] %= MOD

        return a

A = A + A
C = [0] * (N+1)
for r in R:
    C[N-r] += 1

conv = Convolution()
res = conv.convolution(A, C)
print(*res[N:2*N])
0