AtCoder Beginner Contest 335（Sponsored by Mynavi）E,F,G問題メモ

AtCoder Beginner Contest 335（Sponsored by Mynavi）

新年一発目。なかなか大変な新年となっているが。

E - Non-Decreasing Colorful Path

• $N$ 頂点 $M$ 辺の連結無向グラフがあり、頂点 $i$ には整数 $A_i$ が書かれている
• 頂点 $1$ から $N$ まで、以下のルールに従って移動する
  • 同じ頂点は2度通らない
  • 今いる頂点より小さい整数の書かれた頂点には移動できない
• 通った頂点に書かれた整数の「種類数」の最大値を求めよ
  • ルールに従って $1$ から $N$ まで移動することができないときは0
• $2 \le N \le 2 \times 10^5$

最短距離を求めたいならDijkstraなどが使えるが、この問題ではいわば「最長距離」を求めたい。

Dijkstraで“最短”距離が求まるのは、$v$ がキューから取り出された時点で、 もう他のルートで $v$ の今の値よりよい値になるルートはないことが確定しているから。

最長距離の場合、「未探索の迂回ルートで後から更新されうる」可能性がある。
最長路だと思って $u→v$ に探索を広げた後、実は $u$ にもっと長く辿り着けるルートが後から見つかりました、となったら、 $v$ 以降も全て探索しなおしとなり、何回も更新が発生してTLEとなる。

ただ、今回の場合、通るルートが広義単調増加という制約から、上手く評価値を設定することで 「キューから取り出された時点で確定している」ようにできる。

Dijkstraで、キューに入れる値を $(A_v, 暫定種類数 \times -1, v)$ として、小さい方から取り出せばよい。

いま、$(A_v,-S_v,v)$ がキューから取り出されたとして、

• もうそれより小さい値 $A_w$ が書かれた頂点がキューに入ることはない
• また、$A_v$ と同じ値が書かれた頂点から暫定種類数 $S_v$ が更新される可能性もない
  • 同じ値の頂点間の移動では暫定種類数は増えないし、キューに残っているのは $S_v$ 以下の暫定種類数を持つ頂点しかない

よって、もう $v$ に関して $S_v$ が更新されることはない。

この評価値でDijkstra（っぽいこと）を行うことで、$v$ がキューから取り出された時点で、それが最長の $S_v$ だと確定できる。

Python3

from heapq import heappop, heappush

n, m = map(int, input().split())
aaa = list(map(int, input().split()))
links = [[] for _ in range(n)]
rev_links = [[] for _ in range(n)]
for _ in range(m):
    u, v = map(int, input().split())
    u -= 1
    v -= 1
    if aaa[u] <= aaa[v]:
        links[u].append(v)
    if aaa[u] >= aaa[v]:
        links[v].append(u)

scores = [0] * n
scores[0] = 1
q = [(aaa[0], -1, 0)]

while q:
    au, su, u = heappop(q)
    su = -su
    if su != scores[u]:
        continue
    for v in links[u]:
        if au == aaa[v]:
            sv = su
        else:
            sv = su + 1
        if scores[v] < sv:
            scores[v] = sv
            heappush(q, (aaa[v], -sv, v))

print(scores[-1])

F - Hop Sugoroku

• $1,2,...,N$ の番号が振られたマスがあり、はじめ、1にコマがある
• 数列 $A_1,...,A_N$ が与えられ、以下の操作を0回以上、繰り返せる
  • コマを、現在止まっているマス $i$ から、好きな正整数 $k$ を決めて $k \times A_i$ マス進める（$i+k\times A_i$ に移動する）
    • このとき、$N$ を超えるような移動はできない
• 操作を終えた後、「コマが止まったことのあるマスの集合」としてあり得る集合の個数を $\mod{998244353}$ で求めよ
• $2 \le N \le 2 \times 10^5$

愚直には、以下のDPで求められる。

• $DP[i]=$ マス $i$ にコマが止まっているような状態の数

$DP[1]=1$ として、$i=1,2,3,...$ について そこから移動可能なマス $j=i+A_i,i+2A_i,...$ に $DP[j]+=DP[i]$ としていけば、最終的な $\sum_i{DP[i]}$ が答えとなる。

$A_i$ が大きい場合は移動先の候補が少ないので素直に更新できるが、$A_i$ が小さいと移動先が多く、TLEとなる。

こういうときは平方分割。

以下の2つのデータに分けて状態数を管理する。

• $DP'[i]=$ 定義は上と同様だが、↓Added に記録される分は含まない値
• $Added[d][m]=d$ で割って $m$ 余るようなマスに、追加で加算される値

$d$ は、上限を約 $\sqrt{N}$（$=L$ とする）あたりにしておく。

$i$ を遷移元とする場合、その時点の $DP'[i]$ と、 $\displaystyle \sum_{d=1}^{L}Added[d][i \% d]$ の合計が、本来の $DP[i]$ となる。

で、遷移先への加算は、

• $A_i$ が $L$ 以下の場合、$Added[A_i][i \% A_i]$ に $DP[i]$ を加算
• $A_i$ が $L$ より大きい場合、愚直に $j=i+A_i,i+2A_i,...$ に対して $DP'[j]$ に $DP[i]$ を加算

とすると、全部で $O(N\sqrt{N})$ で答えが求められる。

Python3

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

THR = 450
added = [[0] * i for i in range(THR)]

dp = [0] * (n + 1)
dp[1] = 1
ans = 0
MOD = 998244353

for i in range(1, n + 1):
    k = dp[i]
    for d in range(1, THR):
        k += added[d][i % d]
    k %= MOD

    ans += k
    a = aaa[i - 1]

    if a < THR:
        added[a][i % a] += k
        added[a][i % a] %= MOD
    else:
        for j in range(i + a, n + 1, a):
            dp[j] += k
            dp[j] %= MOD

ans %= MOD
print(ans)

G - Discrete Logarithm Problems

• $N$ 個の正整数 $A_1,...,A_N$ と、素数 $P$ が与えられる
• このうち、以下の条件を満たす添字のペア $(i,j)$ の個数を求めよ
  • $A_i^k \equiv A_j \mod{P}$ となる正整数 $k$ が存在する
  • ※ $i=j$ であるものも数える
• $2 \le N \le 2 \times 10^5$
• $P \lt 10^{13}$

原始根を利用すると、累乗がすごろくになる。

• 原始根の定義と具体例（高校生向け） | 高校数学の美しい物語

P = 13,  原始根 r = 2

      e   0  1  2  3  4  5  6  7  8  9 10 11  |  12 13 ...
2^e % P   1  2  4  8  3  6 12 11  9  5 10  7  |   1  2 ...

13に対して2は原始根の1つなので、$2^0,2^1,2^2,...,2^{11} \mod{13}$ は $1～12$ を一巡する。

この上で $A_i=3$ を考えると、$2^e=3$ となるのは $e=4$ なので、$3^k$ は上記の表を4ずつ進むループ状のすごろくとなる。

      e   0  1  2  3  4  5  6  7  8  9 10 11  |  12 13 ... 16   ...
2^e % P   1  2  4  8  3  6 12 11  9  5 10  7  |   1  2 ...  3
         3^0         3^1         3^2             3^3       3^4  ...
↓↓↓
      k   0  1  2  3  4 ...
3^k % P   1  3  9  1  3 ...

もう一例、$A_i=5$ を考えると、$e=9$ なので、$5^k$ は上記の表を9ずつ進む。

      e   0  1  2  3  4  5  6  7  8  9 10 11  |  12 13 ... 18  | ... 27  ... | 36  ...
2^e % P   1  2  4  8  3  6 12 11  9  5 10  7  |   1  2 ... 12  | ...  8  ... |  1  ...
         5^0                        5^1                    5^2       5^3       5^4 ...
↓↓↓
      k   0  1  2  3  4  5  ...
5^k % P   1  5 12  8  1  5  ...

すごろくは $P-1=12$ 周期なので、$A_i=2^e$ とした時、 $A_i^k$ は、$GCD(e,12)$ の倍数のマスは全て止まり、それ以外は止まらない、ということになる。

      e      0  1  2  3  4  5  6  7  8  9 10 11  |  12 13 ...
r^e % P      1  2  4  8  3  6 12 11  9  5 10  7  |   1  2 ...
3^k が止まる *           *           *               *          e=4、GCD(4,12)=4
5^k が止まる *        *        *        *            *          e=9、GCD(9,12)=3

なので、$A_i=2^{e1},A_j=2^{e2}$ として、「$GCD(e_1,12)$ が $GCD(e_2,12)$ の約数なら、$e_1$ ずつ進むすごろくは $e_2$ のマスに止まる」ということになる。
これを言い換えると、元の条件である「$A_i^k \equiv A_j \mod{13}$ となる $k$ が存在する」ことになる。

各 $A_i$ に対して $e_i$、ひいては $GCD(e_i,P-1)$ がわかればよい。
これは $P-1$ の約数なので、種類数は限られる。
$GCD(e,P-1)$ ごとに個数をカウントしておくと、種類数を $L$ として、$O(L^2)$ かけて組合せを総当たりして間に合う。

とはいえ、まず $P$ の原始根を求めるのってそれなりに計算量がかかるし、さらに $A_i$ がその何乗か $e$ を求めるのも大変。

なので、原始根や $e$ は明示的には求めないで、直接 $GCD(e,P-1)$ を求める工夫が必要となる。

位数を求める。

• 位数の性質と原始根の応用 | 高校数学の美しい物語

$a^0,a^1,...$ がどんな周期で“1”に戻るかを示す値だが、必ず $P-1$ の約数となる。

たとえば $P=28001$ として、$P-1$ を素因数分解し、$2^5 5^3 7^1$ となったとする。
$A_i=20$ の位数を求めるとする。以下、mod P は省略する。

20^28000 は必ず 1 （フェルマーの小定理）

2の素因数について
  20^14000 が1か調べる → 1
  20^ 7000 が1か調べる → 1
  20^ 3500 が1か調べる → 1
  20^ 1750 が1か調べる → 1
  20^  875 が1か調べる → not 1
  
  → 長くとも 1750 周期で1になる

5の素因数について
  20^  350 が1か調べる → not 1
  → 5については周期を更新できず、1750周期のまま

7の素因数について
  20^  250 が1か調べる → not 1

→ 20の位数は 1750 とわかる

で、$GCD(e,P-1) = \dfrac{P-1}{位数}$ となるので、つまり $A_i=20$ に対する $GCD(e,P-1)$ は「$28000/1750=16$」となる。

1つの $A_i$ につき、$P-1$ の素因数の個数は $O(\log{P})$、累乗を計算するのに $O(\log{P})$ かかるので、$O((\log{P})^2)$ かかる。

時間制限は長めなので、これで間に合う。

Python3

from collections import Counter, defaultdict


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


n, p = map(int, input().split())
m = p - 1

aaa = Counter(map(int, input().split()))
can = {a: m for a in aaa.keys()}

pf = prime_factorize(m)
for pr, c in pf.items():
    new_can = {}
    pc = pow(pr, c)
    for a, k in can.items():
        k //= pc
        x = pow(a, k, p)
        while x != 1:
            k *= pr
            x = pow(x, pr, p)
        new_can[a] = k
    can = new_can

orders = defaultdict(int)
for a, k in can.items():
    orders[m // k] += aaa[a]

sorted_orders = sorted(orders.keys())
l = len(sorted_orders)

ans = 0
for i in range(l):
    o1 = sorted_orders[i]
    c1 = orders[o1]
    ans += c1 * c1
    for j in range(i + 1, l):
        o2 = sorted_orders[j]
        if o2 % o1 == 0:
            ans += c1 * orders[o2]

print(ans)