ABC162 E - Sum of gcd of Tuples (Hard)

コンテスト中に解けなかった！悔しい！
E - Sum of gcd of Tuples (Hard)

自分の解法や思考過程などを書いています.

解法

競プロフレンズさんのツイートにもあるように問題を読み変え, gcdが $i (1 \le i \le K)$ になるような数列を数え上げる. $i \times (gcdがiになる数列の数)$ の $i (1 \le i \le K)$ で合計を取ったものが答えである.
ただこの問題文を読み変えるのも時間がかかった…(思考メモは下の方に書いた). それにgcdがiになる数列を数え上げるのも, コンテスト中にはできなかった.

ではgcdがiになる数列を数える. どういう数列ならばgcdがiになるか？を考えると, iの倍数のみからなる数列を数え上げることをまず思いついた.
例えばN=3,K=4のとき, gcdが2になるのは(2,2,2)や(2,4,2),(4,2,4)などだ. それぞれの組が2の倍数しか含んでいないことに目をつけてほしい. iの倍数は1からKまでの整数のうちでK//i個含まれるから, iの倍数のみからなる数列の数は, pow(K//i, N, MOD)個である. powの説明はココ.

じゃあgcdがiになる数列を数え上げられたか, というとそうではない. 実は, それらの数列それぞれのgcdが, すべてiにならないからだ. 例を通して説明する.
さっきのN=3,K=4のときに2の倍数のみからなる数列の数は, pow(4//2,3,MOD) = (4//2)³ = 2³ = 8個だ. ただ, gcdが2である数列を全列挙してみる(紙に書いて実験). すると, (2,2,2), (2,2,4), (2,4,2), (4,2,2), (2,4,4,) ,(4,2,4), (4,4,2) のように7個しかない. 何かをよけいに数えてしまったのだ.
実は(2の倍数,ここでは2か4を選ぶか)³をすると, すべて4を選んだときの(4,4,4)の組を含み, このときのgcdは4になってしまう(あ～あ).
iの倍数からなる数列のうち, gcdがiにならないものを除きたい.
~~ここからコンテスト中に詰まった、いやこんなにきれいに考察していなかったけれど~~

どうするか？結論から言うと, gcdがiの倍数になる数列の数から, gcdがiを除くiの倍数になる数列の数を引いていけばいい. そうするとgcdがiになる数列の数が求まる.
$(gcdがiになる数列の数) = (gcdがiの倍数になる数列の数) - (gcdが2 \times i,3 \times i,\cdotsになる数列の数)$ 上の例で見ると,
$(gcdが2になる数列の数) = (gcdが2の倍数になる数列の数) - (gcdが4になる数列の数)$
をしたい. gcdが4になる数列の数を先に求める必要があって, これは1個である(さっき列挙した通り(4,4,4)しかない. pow(4/4, 3)をした後に, gcdが4より大きい4の倍数になる数列の数を引くことを考えてもいいが, K=4なので存在しない). よって, gcdが2になる数列の数は $8 - 1 = 7$ 個になる.

このように, gcdがiになる数列の数を, 上から求めていって(iをKから1までforなどで回す), gcdがiになる数列の数を全て求めれば, 最初に書いたように問題の答えが求まる.

提出コード

実装の注意ですが, PyPyで提出している人は, MODをmainの外に書きましょう.
早くなります. ソースはsolzardさんのツイートです.
PyPy7.3.1からリスト内包表記も早くなるみたいですが, PyPy7.3.0だとリスト内包表記がちょっと遅いみたい(？).
ここまで読んでいただきありがとうございました(以下, コンテスト中の思考過程のメモなので読まなくていいです. 気になる人だけどうぞ).

思考メモ

問題を読んで, まず思いついたのはgcdを何度も計算するのは無駄そうなのでメモして計算をまとめあげるのかな？という発想. 木dpっぽくまとめられたりしないかな？と考えたりしたが, 数列がK^N個もあると計算量は大きくなってしまいそうなので方針転換. ここまで2分くらい？計ってないからわからないけれど.

とりあえずN=2,K=3を書きだして, 見てみる. 一つでも数列に1が含まれてたらその数列のgcdは1, もっと言えば互いに素な組が一つでもあったらgcdは1かと気づく. しかも1,2と2,1のように別の数なら2回数え上げてるな. うーんこれ上手く使えないかな. で, 思いついたのは(互いに素でないグループの大きさ)*(何か)+(互いに素なグループの大きさ)とかにすればいいか？のようなこと. 後者からまず考えたいな, 前者はとりあえず放置. 互いに素も置いておいて, 1を1からN個含むパターンを考えてみる. all 1は一通り, N-1個1はK通り, N-2個1はK-1通り… あ, $\displaystyle{\sum_{i=0}^{N-1} K^i}$ くらいある？あーでも1が一つとか考えた時に思ったけど, 1の場所もかかわってくるか？うーん順列かけてなんとかする？うーーんよくないな

よくわからんので後者の方針を見に行こう互いに素でないグループの和は…例えば偶数のみを考えた時、2がgcdになるな、3がgcdになるのは3の倍数のグループのとき…あれ？倍数に注目すればいいのでは？となる. iの倍数のグループを考えて、大きさが(K//i)なら iは一個含みたい, (K//i)^N-1とN_C_1(これはiを入れる場所)かければいい？この1個特別扱いすればいいみたいなの実装がキレイにならないし, 包除っぽく感じたらとりあえず含んで後で取り除くみたいな気持ちになりたいな. ここらへんが数え上げの感覚がないのか, なにか魔法があると思い込んでたのか詰まってしまったで、もう実装し始めちゃったけれどこれもあまりよくなかったような気がする. もっと考察詰めるべきだった, 焦って実装に走らないは教訓. それから全然答えが合わず眺めていたら, 「gcdがiになるグループ」を数え上げて, グループの大きさ×iの和を取ればいいんじゃね？となったあたりでコンテスト終了. 終了後, TLからヒントをもらってなんとかAC. 次のコンテストではEを解きたい！

翌日, かつっぱさんの質問箱で解法ガチャプロセスを作りましょう, というのを見た. コンテスト中に解法にたどり着けるように, 考察をきれいに早くしたい.