yukicoder No.3257 +|+の別解
2025/09/07
トップページ前の関連記事: yukicoder No.2129 Perfect Binary Tree...?の別解
twitter pixiv yukicoder お題箱 マシュマロ
今回扱うのはyukicoder contest 481(菁々祭プログラミングコンテスト2025)F問題 No.3257 +|+です。公式解説の補足として、相加相乗平均を使うテクニックを解説します。
問題の概要
正確な問題文はこちらです。
長さ$N$の正整数列$A = (A_i)_{i=1}^{N}$が与えられる。$1 \leq i < j \leq N$かつ$i+j \mid A_i + A_j$を満たす整数の組$(i,j)$の個数を求めよ。
前置き
「割り切れる」という関係は「割った余りが$0$」という関係に翻訳でき、割った余りに関する数え上げ問題や総和計算問題は商のfloorの問題に帰着させることが多いです。
何故ならば、商のfloorは単射でなく像が小さいので、同じ値で纏め上げる広い意味での分割統治法が有力だからです。
公式解説ではfloorを取らない商 \[ \frac{A_i + A_j}{i+j} \] を$d$と置き、$d$ごとに計算を行う分割統治法を行っています。ですが商が整数値になる組$(i,j)$の数え上げ問題なので、ここでfloorを取っているか否かは大きな違いではありません。
ただし$d$ごとに単純に計算を行うだけだと$\Theta(N \max(A))$の計算時間がかかるため、公式解説では探索範囲を絞り込む考察と調和数列の増大度の考察を組み合わせて$O(N \log_2 \max(A))$に高速化しています。
ここではもう少し考察の楽な解法として、相加相乗平均を用いる解法を紹介いたします。調和数列への帰着がすんなり思いつく人には不要な解説となりますが、汎用的な計算量評価のテクニックのストックは多く持っておくに越したことはないと思いますので。
しかしながら計算時間は$\Theta((N \max(A))^{2/3})$と重めなので、高速な言語限定の解法だと思ってください。
解説
商のfloorのように単射でない関数(像が小さい関数)の値で纏め上げる分割統治法で相加相乗平均の関係を用いて高速化するテクニックが★3典型です。
相加相乗平均の関係を用いる際はまず後で値を決める正整数変数$B$(バケットサイズ)を取り、纏め上げる値が$B$以上か$B$未満かで場合分けをして処理します。(「$B$以下」と「$B$より大きい」で分けてもよいですが、floorは下からの等号付き不等号と相性が良いので「$B$以上」と「$B$未満」で分ける方が考察が楽です)
本問は商のfloorではなく商に関する問題ですが、商が整数値であるという条件を扱う上ではfloorの議論にある程度帰着できるので、商のfloorの値で纏め上げます。 \[ d(i,j) = \frac{A_i + A_j}{i+j} \] と置くと、不等式$B \leq d(i,j)$は \[ B \leq \lfloor d(i,j) \rfloor \] および \[ i + j \leq \frac{A_i + A_j}{B} \] のいずれとも同値であり、2個目の関係式の右辺は$2 \max(A)/B$以下なのでこの関係式を満たす組$(i,j)$は$O(\min(N,\max(A)/B)^2)$個です。
$B \leq d(i,j)$を満たす組$(i,j)$は愚直に全探索することにして、$d(i,j) < B$を満たす組$(i,j)$に絞った数え上げを考えます。
まず$0 \leq d < B$を満たす整数$d$を決め打って、$d = d(i,j)$を満たす組$(i,j)$の個数を計算することにします。
公式解説と同様に数列$B_d = (B_{d,i})_{i=1}^{N}$を$(A_i-di)_{i=1}^{N}$と定めます。実装上は新たな配列を構成せずとも$d$を小さい順に走査して$A$をinplaceに更新していけば良いです。
$d = d(i,j)$は$B_{d,i} + B_{d,j} = 0$と同値です。一般に与えられた二元一次方程式$sx+ty=u$($s,t,u$は定数)と配列$C = (C_i)_{i=1}^{N}$の合成で得られる関係式$sC_i + tC_j = u$を満たす組$(i,j)$の数え上げ問題は頻度表を用いた走査で$O(N)$で計算できるという典型(いわゆるZero-Sum Ranges)があります。
今回は$(s,t,u,C) = (1,1,0,B_d)$という状況ですし、AtCoder Grand Contest 023 A - Zero-Sum Rangesは$(s,t,u,C) = (1,-1,0,(\sum_{j=1}^{i} A_j)_{i=0}^{N})$という状況ですね。
頻度表はバケットソートで実装できるサイズですし、今回は手抜きで連想配列(unordered_map)で実装してもc++ならば間に合いました。
以上より、この解法の計算量は$O(\min(N,\max(A)/B)^2+BN)$であり、特に$O((\max(A)/B)^2+BN)$です。$O$の中身を最小化すると、$B$を$\sqrt[3]{\max(A)/N}$付近に選ぶことで$O((N \max(A))^{2/3})$となります。
では実際にどのように$O$の中身を最小化すればよいかについて、次で説明していきます。
ちなみに別に最小化しなくても$B$を$\sqrt{N}$付近に取れば$O(\min(N,\max(A)^2/N)+N^{3/2})$になってどうせ通るだろうと思ったら実装が下手なせいか普通にTLEしてしまいました。unordered_mapを過信しすぎたかもしれません。最小化しないのであれば頑張ってバケットソートを書いたほうが良さそうですね。
相加相乗平均
$f(B) = sB^{-n} + tB^m$($s,t,n,m$は正実数定数、$B > 0$)の形の関数は限界費用曲線のようなグラフになります。
この最小化問題は高校数学だと微分を用いる方法が一般的ですね。
\[ f’(B) = -nsB^{-n-1} + mt B^{m-1} = mt B^{-n-1} \left( B^{n+m} - \frac{ns}{mt} \right) \]
の唯一の零点
\[ B = \left( \frac{ns}{mt} \right)^{\frac{1}{n+m}} \]
付近の整数値を$B$に代入するのが正解です。
とはいえ微分はそこそこ面倒なので、もう少し初等的な計算で求めたいです。
そこで$n,m$が小さい整数の時は相加相乗平均の関係を使うとすんなり求められることを覚えておきましょう。
例えば$(n,m)=(1,1)$の時は
\[ f(B) = sB^{-1} + tB \geq 2 \sqrt{sB^{-1} \times tB} = 2 \sqrt{st} \in \Theta(\sqrt{st}) \]
で等号成立条件は$sB^{-1} = tB$です。つまり定数倍を無視すれば$B$を$\sqrt{s/t}$付近に取ればいいですね。
元の問題は$(n,m) = (2,1)$でしたね。この時は
\[ f(B) = sB^{-2} + tB = sB^{-2} + 2^{-1}tB + 2^{-1}tB \geq 3 \sqrt[3]{sB^{-2} \times 2^{-1}tB \times 2^{-1}tB} = 3 \sqrt[3]{st^2} \in \Theta(\sqrt[3]{st^2}) \]
で等号成立条件は$sB^{-2} = 2^{-1}tB$です。つまり定数倍を無視すれば$B$を$\sqrt[3]{s/t}$付近に取ればいいですね。
$(n,m) = (1,2)$の時は
\[ f(B) = sB^{-1} + tB^2 = 2^{-1}sB^{-1} + 2^{-1}sB^{-1} + tB^2 \geq 3 \sqrt[3]{2^{-1}sB^{-1} \times 2^{-1}sB^{-1} \times tB^2} = 3 \sqrt[3]{s^2t} \in \Theta(\sqrt[3]{s^2t}) \]
で等号成立条件は$2^{-1}sB^{-1} = tB^2$です。つまり定数倍を無視すれば$B$を$\sqrt[3]{s/t}$付近に取ればいいですね。
一般の正整数$(n,m)$に対しても
\[ \begin{array}{rcl} f(B) & = & sB^{-n} + tB^m = \underbrace{m^{-1}sB^{-n} + \cdots + m^{-1}sB^{-n}}_m + \underbrace{n^{-1}tB^m + \cdots + n^{-1}tB^m}_n \\
& \geq & (n+m) \sqrt[n+m]{(m^{-1}sB^{-n})^m \times (n^{-1}tB^m)^n} = (n+m) \sqrt[n+m]{(s/m)^m (t/n)^n} \\
& \in & \Theta((n+m)\sqrt[n+m]{(s/m)^m (t/n)^n}) \end{array} \]
で等号成立条件は$m^{-1}sB^{-n} = n^{-1}tB^m$です。つまり$B$を
\[ \left( \frac{ns}{mt} \right)^{\frac{1}{n+m}} \]
付近に取ればいいですね。微分で求めた極と同じです。特に$n,m$が小さい定数の時は$B$を
\[ \left( \frac{s}{t} \right)^{\frac{1}{n+m}} \]
付近に取ればいいというわけです。
実装例
自分用覚え書き
github pagesで数式を使う際は、閉じカッコと半角アンダーバーの間に半角スラッシュを入れないとmarkdownからhtmlへの翻訳の仕様でバグるらしい。
数式環境をくくる大括弧と数式環境内の中括弧や半角空白や濃度のシャープは半角スラッシュ2つ、改行は半角スラッシュ5つ、その他は基本的に半角スラッシュ1つで良さそう。htmlとmathjaxの関係が難しい。
inline数式内の縦棒はlvertなどを使わないとmarkdownの表と認識される。
arrayとalignはどちらも使える。arrayが使えない気がしてしまったのは中括弧につける半角スラッシュの個数を間違えていたから。alignはアンパサンドを等号の右側につけて半角空白を2つ入れると幅がちょうど良さそう。arrayなしだとアンパサンドは使えない。