AtCoder Beginner Contest 230 F,G問題メモ

AtCoder Beginner Contest 230

平日コン、のーみそが疲れてると問題文の内容が頭に入ってこない。

F - Predilection

• 長さ $N$ の数列 $A_1～A_N$
• 隣り合う2要素を、その和の1要素に置き換える、という操作を何度でも行える
• 最終的な数列としてあり得るものは何通りか、$\mod{998244353}$ で求めよ
• $1 \le N \le 2 \times 10^5$
• $|A_i| \le 10^9$

累積和で考える。

$A$ の累積和を $B$ 、最終的な数列の一例を $C$ 、その累積和を $D$ とする。

すると $D$ は、$B$ からいくつかを（順番を変えずに）かいつまんだ部分列と見ることができる。
（ただし最後の $B_N$ は必ず使う）

A   1  0  2  1 -3  2  1  0
B   1  1  3  4  1  3  4  4

    |-->  >  |----------->
C      1  2              1
D      1  3              4

$C$ と $D$ は1対1対応するので、$D$ の個数を数えればよい。

列として同じになってしまうものが存在するので、それを重複して数えてはいけない。

i   1  2  3  4  5  6  7  8
--------------------------
A   1  0  2  1 -3  2  1  0
B   1  1  3  4  1  3  4  4
    >  |----------->  |-->
C   1              2     1    上の例とC,Dが同じ
D   1              3     4

DPで考える。

• $DP[i]=i$ 番目の要素まで考えて、$i$ 番目の要素は必ず使うときの、$B$ から取れる部分列の個数

遷移は、基本的には $DP[1]～DP[i-1]$ のそれぞれから遷移できて以下のようになる。

• $\displaystyle \sum_{j=1}^{i-1}DP[j]$

ただし、複数含まれる要素は最後の1要素からしか遷移できない。

i   1  2  3  4  5  6  7  8
--------------------------
B   1  1  3  4  1  3  4  4
                      *       DP[7] へ遷移できるのは
             *  *  *          この3つだけ

i=3 と i=6 は Bi が 3 で同じ。
DP[3] に数えられている具体的な部分列は (1 3) (1 1 3) などがあるが、
これらは必ず全て、DP[6]にも数えられている。

これは、$i=1,2,...$ と $DP$ を求めていく過程で、「要素 $a$ が最後に出現した $i$」を辞書型などで更新していくとよい。

$O(N)$ で最後まで求められる。

Python3

from itertools import accumulate

n = int(input())
aaa = list(map(int, input().split()))
MOD = 998244353

acc = list(accumulate(aaa))

last = {}
dp = [0] * n
tmp = 1

for i, a in enumerate(acc):
    dp[i] = tmp
    tmp *= 2
    if a in last_i:
        tmp -= dp[last[a]]
    tmp %= MOD
    last[a] = i

print(dp[-1])

G - GCD Permutation

• $1,..,N$ の順列 $P_1,...,P_N$ が与えられる
• 以下を全て満たす $(i,j)$ の組の個数を $\mod{998244353}$ で求めよ
  • $1 \le i \le j \le N$
  • $GCD(i,j) \neq 1$
  • $GCD(P_i,P_j) \neq 1$
• $1 \le N \le 2 \times 10^5$
• 時間制限: 5sec.

基本的に公式Editorialの踏襲。

仮に、$GCD(P_i,P_j) \neq 1$ を無視して $GCD(i,j) \neq 1$ だけに絞って考えると、

• 互いに2の倍数だったら当てはまる
• 互いに3の倍数だったら当てはまる
• ……

と、互いに何らかの数の倍数だったら当てはまる。

たとえば2の倍数は $\dfrac{N}{2}$ 個（切り捨て）あるので、 これを $s$ とおくと $\dfrac{s(s+1)}{2}$ 個のペアが存在する。
（この問題では自分自身とのペアも許される点に注意）

$a=2～N$ について、$a$ の倍数の中で作れるペアの個数 $Pair(a)=\dfrac{\frac{N}{a}(\frac{N}{a}+1)}{2}$ を合計すると答え……にはならない。

重複が発生するので、上手く除かないといけない。
たとえば6と12は、$2,3,6$ の3回で数えられてしまう。

そういうときに便利に使えるのがメビウス関数（を改変したもの）となる。
包除原理の約数特化版みたいな感じ。

• 本来の メビウス関数 は「互いに素なもの」の個数を数える時などに使うため、下記と少し異なる

\[ \tilde{\mu}(n)= \begin{cases} 1 & (nが相異なる奇数個の素数の積で表されるとき (n=2,3,30 など) ) \\ -1 & (nが相異なる偶数個 (0 を含まない) の素数の積で表されるとき (n=6,10,210 など) ) \\ 0 & (それ以外のとき (n=1,4,8,9,12 など) ) \end{cases} \]

2以上のどんな整数についても、約数の $\tilde{\mu}(x)$ を足し合わせると、必ず $1$ になるという性質を持つ。

なので、6と12の場合も、

• $\tilde{\mu}(2),\tilde{\mu}(3)$ は1なので、その倍数から作られるペアは $1$ 回ずつ数えられる
• $\tilde{\mu}(6)=-1$ なので、その倍数から作られるペアは $-1$ 回数えられる（負に寄与する）

とすると、きちんと重複をのぞける。

よって、$1～N$ で互いに素でない組の個数は、$\displaystyle \sum_{a=2}^{N}\tilde{\mu}(a)Pair(a)$ となる。

さて、これでGCD条件が1つの場合は答えを求められた。
しかし今回はGCD条件が2つある。

$a=2～N$ について、添字が互いに $a$ の倍数である $P_a,P_{2a},P_{3a},...$ の中に限定して、 再び $b=2～N$ について「2の倍数の個数」……「$N$ の倍数の個数」を調べていけば答えは求まる。

$g(a,b)$ を「添字が $a$ の倍数である $P_i$ に限定した中での $b$ の倍数の個数」とすると、 $\displaystyle \sum_a \sum_b \tilde{\mu}(a)\tilde{\mu}(b)\dfrac{g(a,b)(g(a,b)+1)}{2}$ が答えとなる。

添字ならば規則的に並んでいるため「$a$ の倍数の個数」は割り算で求められたが、 抽出された中では「$b$ の倍数の個数」は実際に全部試すしかない。 そのためこれを愚直に行うと、計算量は全体で $O(N^2)$ となってしまう。

しかし、個数が限定されている場合、メビウス関数の性質も併せて考えると、実は調べるべき $b$ はそれほど多くない。

$b$ の方を1つ1つ試すのではなく、各 $P_a,P_{2a},...$ について、それらが持つ約数の方からアプローチする。

つまり、「$count[b]=b$ を約数に持つ、限定された中での $P_i$ の個数」とした記録用辞書を作っておき、
たとえば $P_i$ が $2,3,6$ を約数に持つのであれば、$count[2],count[3],count[6]$ にそれぞれ+1する。

$P_a,P_{2a},...$ についてこれを行えば、$count$ から各値を約数に持つ $P_i$ が何個あるかがわかる。

だがそれでも、約数は $N \le 2 \times 10^5$ の中では最大 $160$ 個とかになり、言語によってはやや厳しい。

さらに絞ると、同じ素因数を2個以上含む約数は $\tilde{\mu}(x)$ が $0$ になる（答えへの寄与も0）ので、 約数のうちそれらは無視できる。
$P_i$ が $2^p3^q5^r$ という風に素因数分解できるなら、 $count$ に記録すべき約数は、$2,3,5$ をそれぞれ含むか含まないか、から $1$ を除いた $2^3-1=7$ 通りだけでよい。

これなら、$P_i \le 2 \times 10^5$ の範囲で素因数の種類数はたかだか $6$ 個（考慮すべき約数は $2^6-1=63$ 個）、 しかもそんな特異な数は限られていて多くはせいぜい $1～3$ 個くらいなので、間に合う。

あらかじめ、$1～N$ について「$\tilde{\mu}(x)$」および「約数のうち $\tilde{\mu}(x)$ が0でないもののリスト」を計算しておく。

$a=2～N$ について、以下の要領で $P_a,P_{2a},...$ を調べていく。

• 空の記録用辞書 $count$ を作る
• $P_a,P_{2a},...$ について、約数リストに存在する数について、$count[b]+=1$ としていく
• 最終的に $count$ に記録された各 $(b, cnt)$ につき、$\tilde{\mu}(a)\tilde{\mu}(b)\dfrac{cnt(cnt+1)}{2}$ を答えに加算する

$1～N$ の素因数の種類数の最大値を $d(N)$ として、$O(2^{d(N)} N \log{N})$（実際はもっと少ない）で求められる。

Python3

from collections import defaultdict


def customized_mobius(n):
    table = list(range(n + 1))
    for i in range(2, int(n ** 0.5) + 1):
        if table[i] == i:
            for j in range(i * i, n + 1, i):
                table[j] = i
            for j in range(i * i, n + 1, i * i):
                table[j] = 0
    res = [0] * (n + 1)
    res[1] = -1
    for i in range(2, n + 1):
        if table[i] == 0:
            continue
        res[i] = -res[i // table[i]]
    res[0] = res[1] = 0
    return res


def sequence_prime_factorize(n):
    """ 個数は求めず、含まれる素因数の集合だけ求める """
    table = list(range(n + 1))
    res = [set() for _ in range(n + 1)]
    for i in range(2, n + 1):
        if table[i] != 1:
            res[i].add(table[i])
        if table[i] == i:
            for j in range(i * 2, n + 1, i):
                while table[j] % i == 0:
                    table[j] //= i
                res[j].add(i)
    return res


n = int(input())
ppp = [0] + list(map(int, input().split()))

mb = customized_mobius(n)
pf = sequence_prime_factorize(n)
dv = [[], []]

for i in range(2, n + 1):
    tmp = [1]
    for p in pf[i]:
        tmp.extend([a * p for a in tmp])
    dv.append(tmp[1:])

ans = 0

if ppp[1] == 1:
    ans += n - 1
else:
    ans += n - 2

for i in range(2, n // 2 + 1):
    if mb[i] == 0:
        continue
    check_divisors = defaultdict(int)
    for j in range(i, n + 1, i):
        for d in dv[ppp[j]]:
            check_divisors[d] += 1
    for d, c in check_divisors.items():
        if c >= 2:
            pairs = c * (c - 1) // 2
            ans += mb[i] * mb[d] * pairs

print(ans)