はじめに

今回はABC471 E問題をShrike-Liteで解いてみます。

最初に問題とその公式解説を眺めたときの印象は、かなり良いものでした。

配列そのものを保存しなくてもよさそうです。
必要そうなのは、その時点で読み終わっている値の合計や二乗和など、ほんの数個の状態だけ。

「これなら1120 LUTのShrike-Liteでもいけるのでは？」

と思いました。

しかし、この浅知恵が後に大変な苦労をひき起こします。

その顛末は先のお楽しみ、ということで、まずは問題の解法をFPGA向けに整理してみます。

問題

整数列

A1, A2, ..., AN

が与えられます。

このN個からちょうどK個を選び、そのK個の和を二乗します。

すべての選び方について、その値を足し合わせたものを

mod 998244353

で求めます。

そのまま全組合せを列挙するのはもちろん無理です。

そこで、まず式を展開してみます。

具体例で式を眺めてみる

いきなり一般式から考えると少し見通しが悪いので、まずは小さな例で考えてみます。

例えば、

N = 5
K = 3
A = (A, B, C, D, E)

とします。

5個の要素から3個を選ぶ方法は、

ABC  ABD  ABE
ACD  ACE  ADE
BCD  BCE  BDE
CDE

の10通りです。

この10通りそれぞれについて、

(選んだ3要素の和)^2

を計算し、全部足せば答えになります。

しかし、問題制約はK <= N <= 20万ですから、すべての組み合わせを列挙して計算して回答、とするわけにはいきません。

そこで、この10個の式にどんな構造があるのかを見てみます。

まず1個だけ展開してみる

例えば ABC を選んだ場合は、

(A + B + C)^2

です。

展開すると、

A^2 + B^2 + C^2
+ 2AB + 2AC + 2BC

となります。

つまり二乗を展開すると、出てくるものは、

各要素の二乗
異なる2要素の積

の2種類に分類できそうです。

少しですが「問題の構造」のようなものが見えてきました。

A^2は何回出てくる？

ここで、すべての組み合わせ10通りをそれぞれ展開する代わりに、

10通りを全部計算して足し合わせると、A^2はいったい何回出てくるのか？

を考えてみます。

10通りの組み合わせのうち、Aを含む3個組は、

ABC
ABD
ABE
ACD
ACE
ADE

の6通りです。

したがって、10通りの式を全部足したとき、A^2 は6回出てきます。

また、Bを含む3個組も、

ABC
ABD
ABE
BCD
BCE
BDE

の6通りですから、B^2 も6回出てきます。

CでもDでもEでも事情は同じです。

したがって、計算全体で出現する二乗項だけを集めると、

6 * (A^2 + B^2 + C^2 + D^2 + E^2)

とその和が計算できるはずですね。

二乗項の登場回数を一般化する

では、この 6 はどこから出てきたのでしょうか。

A^2 が登場するためには、選んだ3個の中にAが含まれていればよいことになります。

そこで、まずAを1個選んでしまいます。

選ぶ3個のうち

1個 : Aで確定
残り: あと2個

A以外には、

B, C, D, E

の4個があります。

この4個から残り2個を選べばよいので、

BC  BD  BE
CD  CE
DE

の6通りです。

組み合わせの記号を使うと、

C(4, 2) = 6

です。

数式では、

\binom{4}{2}=6

とも書きます。意味を日本語で言えば「4個から2個を選ぶ組み合わせの数」です。

この先は、この表記も使っていきます。

では、これを一般の N, K に置き換えてみます。

全部でN個の要素からK個を選ぶとき、ある要素 Ai の二乗 Ai^2 が登場する回数を考えます。

まず Ai を1個選んでしまえば、

残っている要素 : N-1個
あと選ぶ要素   : K-1個

です。

したがって、その選び方は、

C(N-1, K-1)

通りです。

つまり、どの Ai^2 も答えの中に C(N-1, K-1) 回ずつ登場します。

したがって、二乗項全体は、

C(N-1, K-1) * (A1^2 + A2^2 + ... + AN^2)

とまとめられます。

数式表記では

\binom{N-1}{K-1}
\sum_{i=1}^{N} A_i^2

です。

ここで Σ は、

A1^2 + A2^2 + ... + AN^2

を短く書いただけです。

ABは何回出てくる？

次に、異なる2要素の積（交差項）を見てみます。

例えば AB が登場するのは、

ABC
ABD
ABE

の3通りです。

そして、

(A + B + C)^2

を展開すると 2AB が出てきたように、二乗展開ではABは毎回 2AB の形で現れます。

したがって、10通りの式を全部足したときのABの寄与は、

3 * 2AB = 6AB

です。

ACでもADでも、ほかの2要素の組でも事情は同じです。

したがって、交差項だけを集めると、

6(AB + AC + AD + AE
    + BC + BD + BE
    + CD + CE
    + DE)

になります。

交差項の登場回数を一般化する

今度は、AB が3回登場した理由を考えてみます。

ABが登場するためには、選んだ3個の中にAとBの両方が含まれていればよいことになります。

そこで、まずAとBを選んでしまいます。

選ぶ3個のうち

2個 : A, Bで確定
残り: あと1個

AとB以外には、

C, D, E

の3個があります。

この3個から残り1個を選べばよいので、

C
D
E

の3通りです。

組み合わせの記号を使うと、

C(3, 1) = 3

です。

これも一般の N, K に置き換えてみます。

全部でN個の要素からK個を選ぶとき、ある2要素 Ai と Aj の積 AiAj が登場する回数を考えます。

まず Ai と Aj の2個を選んでしまえば、

残っている要素 : N-2個
あと選ぶ要素   : K-2個

です。

したがって、その選び方は、

C(N-2, K-2)

通りです。

さらに二乗展開では、AiAj は毎回、

2 * Ai * Aj

の形で現れます。

したがって、ある AiAj が答えに与える寄与は、

C(N-2, K-2) * 2 * Ai * Aj

です。

すべての異なる2要素の組についてまとめると、

2 * C(N-2, K-2)
  * (A1A2 + A1A3 + ... + A(N-1)AN)

です。

数式表記だと

2\binom{N-2}{K-2}
\sum_{1\le i<j\le N} A_iA_j

です。

ここでも Σ は、「異なる2要素の積を全部足す」という長い足し算を短く書いただけです。

この例の答えはこう書ける

以上から、N=5, K=3 の場合は、

6(A^2 + B^2 + C^2 + D^2 + E^2)

+ 6(AB + AC + AD + AE
     + BC + BD + BE
     + CD + CE
     + DE)

を計算すればよいことが分かります。

ここで、

square_sum =
    A^2 + B^2 + C^2 + D^2 + E^2

pair_sum =
    AB + AC + AD + AE
       + BC + BD + BE
            + CD + CE
                 + DE

と名前を付けると、

6 * square_sum + 6 * pair_sum

です。

かなり構造が見えてきましたね。

あとは具体的にどうやってこれらの数を計算していくかを考えればよさそうです。

pair_sumも配列を保存せずに作れる

square_sumの計算は簡単そうです。

配列の要素を受け取るたびに、その要素の二乗を足していけば、常に最新情報を保持できます。

厄介そうなのはpair_sumです。

そこで

AB + AC + AD + AE
   + BC + BD + BE
        + CD + CE
             + DE

をどう計算していけばいいか、を考えます。

入力が、

A, B, C, D, E

の順に来るとします。

Aが来た時点では、まだ相手がいません。

次にBが来た時点で、Aは知っているから、ABは計算できるようになりますね。

そこで、pair_sumに

AB

を加算します。

Cが来たら、AもBも知ってるから、新たにACとBCが計算できるようになります。

AC + BC
= (A + B)C

を追加します。

同様にDが来たら、

AD + BD + CD
= (A + B + C)D

を追加します。

最後にEが来たら、

AE + BE + CE + DE
= (A + B + C + D)E

を追加します。

ここでこの足し算チェーンをよく見ると、 それ以前に受け取った要素の合計に、新しい要素をかけて足すを繰り返しています。

つまり、Eが到着した時点で必要な情報は(A + B + C + D)の値だけで、AとかBそれぞれの個別の値はなくても大丈夫、ということになります。

したがって、pair_sumの更新は

prefix_sum = これまでに読んだ値の合計
pair_sum   = これまでに完成した2要素積の合計

を持っておけば、新しい値Xが来るたびに、

pair_sum   += prefix_sum * X
prefix_sum += X

と更新していけばいいはずですね。

二乗和の部分も、

square_sum += X * X

と新しい要素を受け取るたびに加算更新していけばいいですから、1要素を受け取るたびに必要な計算を済ませたら、もうその要素は忘れていいことになります。

これなら入力を左から右へ1回読みながら、

prefix_sum
pair_sum
square_sum

の3個を更新し続けるだけで済みます。

公式解説の裏ワザ：さらに保持するデータを減らす

ここまでで、

prefix_sum
pair_sum
square_sum

の3個を更新し続ければよいことが分かりました。

これだけでも十分に小さそうですが、公式解説にはさらに保持するデータを1個減らす裏ワザが紹介されています。

ポイントは、この3個の値がバラバラに変化しているわけではなく、入力をどこまで読んだ時点でも、常に一定の関係を保っていることです。

まずAだけを読み終えた直後は、

prefix_sum = A
square_sum = A^2
pair_sum   = 0

です。

したがって、

prefix_sum^2
= square_sum + 2 * pair_sum

が成り立ちます。

次にBまで読み終えると、

prefix_sum = A + B
square_sum = A^2 + B^2
pair_sum   = AB

です。

このときも、

prefix_sum^2
= (A + B)^2
= A^2 + B^2 + 2AB
= square_sum + 2 * pair_sum

となります。

Cまで読み終えた場合も、

prefix_sum = A + B + C

square_sum =
    A^2 + B^2 + C^2

pair_sum =
    AB + AC + BC

ですから、

prefix_sum^2
= (A + B + C)^2

= A^2 + B^2 + C^2
  + 2(AB + AC + BC)

= square_sum + 2 * pair_sum

となります。

つまりAまで、Bまで、Cまで……と入力が進んでも、常に、

prefix_sum^2
= square_sum + 2 * pair_sum

という関係が保たれています。

したがって、

2 * pair_sum
= prefix_sum^2 - square_sum

です。

今回の答えを計算するときに必要なのは pair_sum そのものではなく、二乗展開で現れる 2 * pair_sum です。

ということは、pair_sum を毎回計算して保持しておかなくても、

prefix_sum
square_sum

の2個だけを更新しておき、最後に

prefix_sum^2 - square_sum

を計算すればよいことになります。

3個でも十分少ないと思っていたのに、2個まで減りました。

これはかなりShrike-Lite向きに見えます。

一般式をまとめる

ここまでで、二乗項と交差項の登場回数をそれぞれ求められました。

二乗項全体は、

C(N-1, K-1)
* (A1^2 + A2^2 + ... + AN^2)

でした。

交差項全体は、

2 * C(N-2, K-2)
* (A1A2 + A1A3 + ... + A(N-1)AN)

でした。

したがって答えは、

Ans =
    C(N-1,K-1) * (A1^2 + A2^2 + ... + AN^2)
  + 2 * C(N-2,K-2)
      * (A1A2 + A1A3 + ... + A(N-1)AN)

となります。

まとめて書くと、

\mathrm{Ans}
=
\binom{N-1}{K-1}
\sum_{i=1}^{N}A_i^2
+
2\binom{N-2}{K-2}
\sum_{1\le i<j\le N}A_iA_j

です。

ここまでの説明で、

square_sum =
    A1^2 + A2^2 + ... + AN^2

そして、

2 * pair_sum
= prefix_sum^2 - square_sum

でした。

したがって、先ほどの長い式は、

Ans =
    C(N-1,K-1) * square_sum
  + C(N-2,K-2)
      * (prefix_sum^2 - square_sum)

とも書けます。

最終的には、

\mathrm{Ans}
=
\binom{N-1}{K-1}\mathrm{square\_sum}
+
\binom{N-2}{K-2}
\left(
\mathrm{prefix\_sum}^2-\mathrm{square\_sum}
\right)

です。

一般式だけを見ると少し大げさですが、やっていることは先ほどの N=5, K=3 と同じです。

意味を日本語で書けば、

各要素の二乗を全部足して、登場回数をかける
異なる2要素の積を全部足す代わりに裏ワザで合計を求めて、それに登場回数をかける
足し合わせて答えを出す

というだけです。

数学の記号は、その説明を短く書くために使っています。

そして入力列を読みながら保持する主要な値は、結局、

prefix_sum
square_sum

の2個だけです。

配列要素は読んだら処理して捨てる、を繰り返しても問題ありません。

これはShrike-Liteでもいけるのでは？

ここまで考えた時点では、かなり楽観的でした。

Shrike-LiteのForge FPGAは大きなFPGAではありません。

しかし今回のアルゴリズムは、

配列を保存しない
巨大な探索状態を持たない
素直な方法でも主要な集計値は3個
入力は一方向

です。

いかにも小規模FPGA向けに見えます。

問題はMODが

998244353

なので、値を保持するには30bit必要なことくらいでしょうか。

公式解説の裏ワザを使えば pair_sum まで消して、集計値は2個だけですむことも分かっています。

ただ、3個でも十分に少ないです。

そこで初期実装では、まず

prefix_sum
pair_sum
square_sum

の3個をそのまま管理する、分かりやすい方法で作ってみることにしました。

「裏ワザまで使わなくても、このくらいなら入るでしょう」

という見通しです。

30bitの演算器は、できるだけ共有することにしました。

初期実装

初期版では、おおむね次の構成にしました。

prefix_sum / pair_sum / square_sum集計
30bit modular add/sub
30bit modular multiplier
組合せ係数計算
Fermatの小定理による逆元計算
最終式の計算
SPI通信

この段階では、先ほど紹介した pair_sum を消す裏ワザはまだ使っていません。

まずは3個の集計値をそのまま持つ、素直な実装です。

乗算器は複数置かず、1個を共有します。

30bitのmodular multiplicationはshift-add型とし、1回の乗算を複数クロックに分けました。

速度より面積を優先したつもりです。

組合せ係数も、まずは

C(n,k)

として考えます。

この分子・分母を順番に計算し、分母の逆元を、

x^(MOD-2) mod MOD

で求めます。

この時点では、

「少し大きいかもしれないけれど、乗算器も共有したし、何とかなるのでは」

くらいに考えていました。

まずIcarusで確認

機能シミュレーションでは問題ありませんでした。

入力を順番に与え、ここまで説明してきた、

prefix_sum
pair_sum
square_sum
組合せ係数
最終結果

に相当する値が期待通りに計算できることを確認しました。

アルゴリズムとしては成立しています。

では、合成してみます。

あれ？

手元に残っている最終BaselineのPNRログでは、packing後の使用量は次のようになりました。

Logic 6-LUT CLB : 577
CLB LUT         : 2190
CLB FF          : 888

Shrike-Liteで使えるLogic 6-LUT CLBは140個です。

つまり、

Type=L : 577 / 140

という世界です。

全然入りません。

「少し大きい」ではありません。

4倍以上あります。

レジスタ3個しかないのに

ここで最初の甘い予想が崩れました。

初期実装で長期間保持している主要な集計値は、

prefix_sum
pair_sum
square_sum

の3個です。

配列全体を保存しているわけでもありません。

しかし、

管理する数学的な状態が少ない

ことと、

回路が小さい

ことは同じではありませんでした。

30bitのmodular arithmetic、組合せ係数、逆元、pow、その制御。

数式では数行で書ける部分が、FPGAではかなり大きな回路になります。

特に、

C(N-1,K-1)
C(N-2,K-2)

を作るための逆元処理が怪しそうです。

普通のプログラムならmodinv一行なのにね......

今回はここまで

最初は、

配列を保存しない。
主要な集計値も3個だけ。
しかも公式解説の裏ワザを使えば2個まで減らせる。
これならShrike-Liteでもいけるのでは？

と思っていました。

そこで初期実装では、

まあ3個でも十分少ないし、裏ワザはまだいいでしょう。

くらいの気持ちで、素直な3個管理のまま作りました。

ところが最終BaselineでPNRを実行すると、配置段階で、

Type=L 577 / 140

となりました。

まったく入りません。

次回は、

何がそんなに面積を使っているのか
どこまで共有できるのか
30bitが悪いのか
組合せ係数が悪いのか

を調べながら、ひたすら狭小住宅に家具を詰め込んでみます。

まずは、先ほど紹介した pair_sum を消す裏ワザから試してみましょう。

前回：
Shrike-LiteでAtCoder問題を解く(36)：ABC471A - 4つの条件を並列判定する

次回：
Shrike-LiteでAtCoder問題を解く(38) ABC471E（後編）の予定

コード全文と実装資料：(コード実装やドキュメント作成は生成AIの助けを借りています)
第37回コード全文と実装資料