Beautiful Palindrome

2 secs 1024 MB

$M = 1$ の場合、 $N ~ \mathrm{mod} ~ 998244353$ が答えです。以降 $M \ge 2$ の場合を考えます。

正整数 $k$ が長さ $M$ の文字列の中に含まれるループ文字の個数としてあり得ることと、 $k - 1$ が $M - 1$ の約数であることは同値です。
理由は、文字列の末尾のループ文字 $1$ 個を除いたとき、 $1$ 個のループ文字とある非負整数個のスペースがこの順に並べられた文字列が $k$ 個並べられた文字列として考えられるためです。

また $k$ 個のループ文字を含む回文の個数は、先頭 $\lceil k / 2 \rceil$ 個のループ文字を決めれば、他の文字も決まるため $N^{\lceil k / 2 \rceil}$ 個となります。
よって、 $\displaystyle \sum_{k - 1 |M - 1} N ^{\lceil k/2 \rceil} ~ \mathrm{mod} ~ 998244353$ が答えです。

$M - 1$ の約数の個数は $O(\sqrt M)$ 個に抑えられ、累乗計算は繰り返し二乗法を用いて $O(\log N)$ で計算可能なため、 $O(\sqrt{M}\log N)$ で解くことができました。

想定解

C++

xxxxxxxxxx
 
#include <bits/stdc++.h>
#include <atcoder/modint>
​
using namespace std;
using namespace atcoder;
​
using mint = modint998244353;
​
vector<int> divisors(int n) {
  vector<int> res;
​
    for (int i = 1; i * i <= n; i++) {
        if (n % i == 0) {
            res.push_back(i);
            if (i * i != n) res.push_back(n / i);
        }
    }
​
    return res;
}
​
int main() {
    int n, m;
    cin >> n >> m;
​
    if (m == 1) {
        cout << n % 998244353 << endl;
        return 0;
    }
​
    mint ans = 0;
​
    for (int div : divisors(m - 1)) {
        ans += mint(n).pow((div + 2) / 2);
    }
​
    cout << ans.val() << endl;
    return 0;
}

Python

xxxxxxxxxx
 
import math
​
MOD = 998244353
​
​
def divisors(n):
    divisors = []
    for i in range(1, int(math.sqrt(n)) + 1):
        if n % i == 0:
            divisors.append(i)
            if i != n // i:
                divisors.append(n // i)
    return divisors
​
​
N, M = map(int, input().split())
​
if M == 1:
    print(N % MOD)
    exit()
​
ans = 0
​
for div in divisors(M - 1):
    ans += pow(N, (div + 2) // 2, MOD)
    ans %= MOD
​
print(ans)
​

Beautiful Palindrome

解説

想定解

C++

Python