AtCoder Regular Contest 167 B,C,D問題メモ

AtCoder Regular Contest 167

B - Product of Divisors

• $A^B$ の正の約数の総積が、$A$ で何回割り切れるか、$\mod{998244353}$ で求めよ
• $2 \le A \le 10^{12}$
• $0 \le B \le 10^{18}$

たとえば $360=2^3 3^2 5^1$ の約数は、(2の次数を0～3から選ぶ)×(3の次数を0～2から選ぶ)×(5の次数を0～1から選ぶ) の $(3+1)(2+1)(1+1)=24$ 個ある。

$A=360,B=5$ とすると $360^5=2^{15}3^{10}5^{5}$ の約数の総積を求めたい。

ひとまず2の次数がどうなるか考える。約数の2の次数 $p$ で場合分けすると、

• $p=0$ の時、3の次数の選び方で11通り、5の次数の選び方で6通り、これらの総積の2の次数は 0
• $p=1$ の時、3の次数の選び方で11通り、5の次数の選び方で6通り、これらの総積の2の次数は 66
• $p=2$ の時、3の次数の選び方で11通り、5の次数の選び方で6通り、これらの総積の2の次数は 132
• …
• $p=15$ の時、3の次数の選び方で11通り、5の次数の選び方で6通り、これらの総積の2の次数は 990

結局、2の次数とは独立に3や5の次数の選び方は決まるので、$p$ 毎に全て同じ個数だけあることになる。

総積の2の次数はこれらの総和なのでまとめることができて、$\dfrac{15 \times 16}{2} \times (10+1)(5+1) = 7920$ となる。
$360^5$ の約数の総積は、$2^{7920}3^{何か}5^{何か}$ だということになる。

もし2の次数だけを考えると（3の次数も5の次数も足りていれば）、 もとの $A=360$ の次数が3だったので、$7920 / 3 = 2640$ 回 割り切れることになる。

で、これを3や5でも求めてみると、どうも総積の素因数分解 $2^x3^y5^z$ の $x:y:z$ の比は、 元の $A$ の素因数分解 $2^a3^b5^c$ の $a:b:c$ と等しくなるっぽい。

なので、どれか1つの素因数だけについて求めればよい。

まとめると、

• $A=p^xq^yr^z$ などと素因数分解する
• $\dfrac{Bx(Bx+1)}{2}(By+1)(Bz+1)$ を計算する（結果を $X$ とする）
• $X$ を $x$ で割ると答え

$x$ についてはあらかじめ割っておけるので、$\dfrac{B(Bx+1)(By+1)(Bz+1)}{2}$ を求めればよいことになる。
この式は $x,y,z$ に対して対称なので、どの素因数でやっても結果は等しいことの証明にも繋がる。

ただし、分子が奇数になり得る。

Pythonでは多倍長整数が使えるのでそのまま計算してもたいした桁数にはならないが、 普通のlong型などを使う場合、$B,Bx+1,By+1,Bz+1$ が全て奇数の場合には「1を引いてから MODINV(2) をかける」のようにする必要がある点に注意。

Python3

from collections import defaultdict


def get_primes(n):
    """ nまでの素数リスト """
    candidates = [0, 1] * (n // 2 + 1)
    if n % 2 == 0:
        candidates.pop()
    candidates[1] = 0
    candidates[2] = 1
    for p in range(3, int(n ** 0.5) + 1, 2):
        if candidates[p]:
            for m in range(p * p, n + 1, 2 * p):
                candidates[m] = 0
    return [p for p, b in enumerate(candidates) if b]


def prime_factorize(x, primes=None):
    """ xを素因数分解、辞書型で返す """
    # primes: sqrt(x)までの素数リスト
    #   省略時は奇数を試し割り（2だけ別処理）
    #   単発なら不要だが、複数のxで行う場合は事前計算した方がよい

    factors = defaultdict(int)
    if primes is None:
        if x & 1 == 0:
            b = (x & -x).bit_length() - 1
            factors[2] = b
            x >>= b
        primes = range(3, x, 2)
    for p in primes:
        if p * p > x:
            break
        if x % p == 0:
            factors[p] = 1
            x //= p
            while x % p == 0:
                factors[p] += 1
                x //= p
    if x != 1:
        factors[x] = 1
    return factors


a, b = map(int, input().split())
MOD = 998244353

primes = get_primes(int(a ** 0.5) + 5)
pf = prime_factorize(a, primes)

pc = list(pf.values())
pcb = pc.copy()
n = len(pc)
for i in range(n):
    pcb[i] *= b

c0 = pcb[0]
pat = c0 * (c0 + 1) // 2
for j in range(1, n):
    c1 = pcb[j]
    pat *= c1 + 1
pat //= pc[0]

print(pat % MOD)

C - MST on Line++

• $N$ 個の整数列 $A=(A_1,...,A_N)$ と正整数 $K$ がある
• $A$ を並べ替えた数列を $B$ とする
  • よりちゃんと定義すると、順列 $P=(P_1,...,P_N)$ を決めて、$B=(A_{P1},A_{P2},...,A_{PN})$ とする
• 以下の問題の答えを $f(B)$ とする
  • $N$ 頂点の無向グラフを考える
    • 2頂点 $(i,j)$ について、$|i-j|$ が $K$ 以下ならコスト $\max(B_i,B_j)$ の辺がある
    • 2頂点 $(i,j)$ について、$|i-j|$ が $K$ より大きいなら辺は無い
  • このグラフでの最小全域木の辺のコストの和を求めよ
• $B$ の決め方は $N!$ 通りあるが、全てに対しての $f(B)$ の総和を $\mod{998244353}$ で答えよ
• $1 \le K \le N \le 5000$
• $1 \le A_i \le 10^9$

辺のコストは $A1,A2,...$ のどれかになるので、 「コストが $A_i$ になるような辺が、$N!$ 通り全体で何本、最小全域木に採用されるか」のようなことを合計できればよい。

いろいろ実験してみると、どうやら、ある $B$ を固定した時の最小全域木は以下のように構築できる（未証明）。

• 各 $i=1～N-1$ に対して、$i$ からは、その右最大 $K$ 個（$j=i+1～\min(N,i+K)$）の中で $B_j$ が最小の頂点につなぎにいく

N=7  K=2
2 7 5 1 4 6 3

以下に辺を張ると最小全域木になる
・2→5
・7→1
・5→1
・1→4
・4→3
・6→3

なので、$i$ ごとに「$A_i$ からつなぎにいく辺のコスト」の全ての順列に対しての総和を計算し、それを全ての $i$ で合計すると答えとなる。

$A_i$ を昇順にソートして、1つ固定すると、

• $A_i$ が $B$ において $1～N-K$ 番目のどこかに置かれたとき、右にある個数は $K$ 個
  • $K$ 個中の最小値が $A_i$ より小さい場合の数 × $A_i$
  • $K$ 個中の最小値が $A_{i+1}$ の場合の数 × $A_{i+1}$
  • $K$ 個中の最小値が $A_{i+2}$ の場合の数 × $A_{i+2}$
  • …
  • $K$ 個中の最小値が $A_{N}$ の場合の数 × $A_{N}$
• $A_i$ が $B$ において $N-K+1$ 番目に置かれたとき、右にある個数は $K-1$ 個
  • 同様に $K-1$ 個中の最小値で場合分けして計算
• $A_i$ が $B$ において $N-K+2$ 番目に置かれたとき、右にある個数は $K-2$ 個
  • 同様に計算
• …
• $A_i$ が $B$ において $N-1$ 番目に置かれたとき、右にある個数は $1$ 個
  • 同様に計算

上記の総和が、$A_i$ に対しての答えとなる。

ただ、このままでは

• 各 $i=1～N$ について、
• 右にある個数別 $k=1～K$
• 右にある中の最小値別 $j=i～N$

にコストを求めているので、計 $O(N^2K)$ となり無理。

$A$ は昇順にソート済みとし、実装を見据えて0-indexedに振りなおす。

以下、順列数え上げ ${}_nP_r = \dfrac{n!}{(n-r)!}$ を用いるが、$n \lt r$ の場合は0とする。

下記の設定を考える。
$i=4$ の50が、$B$ において右に $k$ 個ある位置（$K$ 個以上ある場合は $k=K$）に置かれた場合、 50からつなぎにいく辺のコストは、相手先と $A_i$ との大小関係により、以下のように①②に分けられる。

N=10  K=3  A=(10,20,30,40,50,60,70,80,90,100)

k=3の例  _ _ _ _ 50 _ _ _ _ _
                    ~~~~~
k=2      _ _ _ _ _ _ _ 50 _ _
                          ~~~

①Aiより小さい要素が含まれる並べ方とコスト:
  ・並べ方の個数
      ・k 個内の並べ方
          全ての並べ方 - Aiより大きい要素のみからなる並べ方 = (N-1)P(k) - (N-i-1)P(k)
      ・Aiとk個以外の並べ方
          (N-k-1)!
  
  ・辺コストは Ai
  
  →i,k を固定した総コストは  Ai * ( (N-1)P(k) - (N-i-1)P(k) ) * (N-k-1)!

②最小値が Aj (j>i) となる並べ方とコスト:
  j = i+1,...,N につき、以下の総和
    ・並べ方の個数
        ・k 個内の並べ方
            Ajの位置の選び方 k  ×  残りをAjより大きい要素で埋める (N-j-1)P(k-1)
        ・Aiとk個以外の並べ方
            (N-k-1)!
    
    ・辺コストは Aj
    
    →j,k を固定した時の総コストは  Aj * k * (N-j-1)P(k-1) * (N-k-1)!

※k=K の場合、①②の各結果は、そのようなB上の位置の個数 N-K 倍する

ここで、②は、足し合わせる範囲以外、$i$ に依存しないのがポイント。

なので、各 $j$ について求めておいて、後ろから累積和を取ることで、 各 $i$ に対しての寄与を事前に一括で求めることができる。

①は $i,k$ に対して $O(NK)$ で求められ、②は $j,k$ に対して $O(NK)$ で求め、累積和を取ればよい。

Python3

D - Good Permutation

• $(1,2,...,N)$ の順列 $(Q_1,Q_2,...,Q_N)$ のうち、以下を「よい順列」とする
  • ある要素 $x$（$1 \le x \le N$）からはじめて、$x$ を $Q_x$ に置き換える操作を好きなだけ繰り返せる
  • このとき、どの $x$ から始めても、$x=1$ にすることができる
• 順列 $P=(P_1,...,P_N)$ が与えられる。よい順列とは限らない
• 2項間のスワップを繰り返して $P$ をよい順列にするとき、その最小手数を $M$ とする
• $M$ 回のスワップで達成できるよい順列は複数考えられるが、その中で辞書順最小のものを求めよ
• 1つの入力につき $T$ ケース与えられる
• $1 \le T \le 10^5$
• $2 \le 1つの入力のNの総和 \le 2 \times 10^5$

辞書順なので前から貪欲に決めていく。上手く状態管理してシミュレーションする。

ループを合成して1つにすればよい。

順列は、$x→P_x→P_{Px}→...$ と辿っていくとどこかで $x$ に戻り、ループになる。

ループは複数に分かれている場合もあるが、全ての要素が1に行き着くということは、「よい順列＝ループ1つの順列」となる。

「異なるループに含まれる2要素をスワップ」すると、必ずループが合成され、2つのループの要素を全て巡る1つのループになる。

これを繰り返して1つのループにすればよい。$M=(初期状態のループ数-1)$ 回繰り返すと、必ずループは1つになる。

$1→6→3→2→1$ のようなループを、$\{1,2,3,6\}$ のループとする。

全てのループを「$1$ の属するループ」に合成することを試みる。

辞書順最小は、先頭から貪欲に、最も小さい要素を持ってこれるなら持ってくる、というように確定していける。

i   1 2 3 4 5 6 7 8 9    a: {1,2,4} のループ
Pi  2 4 5 1 3 6 9 8 7    b: {3,5} のループ
    a a b a b c d e d    c: {6} のループ
                         d: {7,9} のループ
                         e: {8} のループ

・i=1, Pi=2
    ・「まだ 1 の属するループ a に合成されてない、Pi より小さい要素」があれば、
      その中の最小値と入れ替えれば、辞書順は最小になる。（"未合成最小値" と呼ぶことにする）
    ・今回の場合、Pi=2より小さい要素は 1 しかなく、1 はループ a に属するので、未合成最小値はない。

・i=2, Pi=4
    ・未合成最小値に 3 がある
    ・4と3を入れ替える
    ・3の属するループ {3,5} が、新たにループ a に加わる

i   1 2 3 4 5 6 7 8 9
Pi  2 3 5 1 4 6 9 8 7
    a a a a a c d e d

・i=3,4,5 … 未合成最小値無し

・ここで、{1,2,...,i} のみからなるループができてしまった（この例では i=5）
    ・この場合、辞書順は悪化するが、大きい要素とスワップしないと、全体が1つにならない
    ・「i+1 である要素 6」と、「i=5（ループ a の一番末尾の要素）である 4」を
      スワップするのが、一番悪化が少ない
    ・6の属するループ {6} が、a に加わる

i   1 2 3 4 5 6 7 8 9
Pi  2 3 5 1 6 4 9 8 7
    a a a a a a d e d

・i=6 ... 未合成最小値無し

・再び、i=6 として 1～i だけでループができてしまった
    ・「i+1 である 7」と、「ループ末尾の 4」をスワップする
    ・7の属するループ {7,9} が a に加わる

i   1 2 3 4 5 6 7 8 9
Pi  2 3 5 1 6 7 9 8 4
    a a a a a a a e a

・i=7, Pi=9
    ・未合成最小値 = 8
    ・9と8を入れ替える
    
i   1 2 3 4 5 6 7 8 9
Pi  2 3 5 1 6 7 8 9 4
    a a a a a a a a a

ループが1つになった！

このように、「未合成最小値を求め、あれば入れ替える」と「$1～k$ ループができていれば末尾と $A_j=i+1$ なる要素を入れ替える」の2つの操作でできる。
それぞれが可能なように状態を管理してシミュレーションすればよい。

• ループ $a$ に属する要素を、合成することも含めて管理する
  • メンバーをsetで保持するUnion-Findなどで可能
  • 合成時、マージテクを使えばマージ部分の計算量は $O(N \log{N})$
• $a$ に属さない中での最小値を得る
  • 未合成最小値は単調に増加していくので、$x=2,3,...$ の順に、その時点の a のメンバーに属するかどうか調べる
  • 見つかったら打ち切って、次は途中から調べれば、全体で $O(N)$
• 「$1～k$ のみからなるループ」かどうか調べる
  • $i$ についての未合成最小値を処理後、$|a|=i$ となっているかどうかで判定できる
• シミュレーションするのに、要素 $P_j$ がどこにあるかが必要
  • $P_j$ から $j$ を逆引きできるリストを管理

などがあればよい。

Python3