トヨタ自動車プログラミングコンテスト2024#12（AtCoder Beginner Contest 384）F,G問題メモ

トヨタ自動車プログラミングコンテスト2024#12（AtCoder Beginner Contest 384）

PythonでTLEになったとき、C++でごり押すために生成AIで変換することは許されてるけど、 結局C++上のエラーを解決する力が自分自身に無ければ、一発で上手くいかない限り時間ロスするだけに終わるんだよなぁ。。。
うまく解釈してもらうためにPython書き直すのもなんか違うし。

F - Double Sum 2

• 正整数 $x$ に対して $f(x)$ を「 $x$ が偶数である間 $x$ を $2$ で割り続けたときの、最終的な $x$ の値」として定義します。例えば $f(4)=f(2)=f(1)=1$ 、 $f(12)=f(6)=f(3)=3$ です。
• 長さ $N$ の整数列 $A=(A_1,A_2,\ldots, A_N)$ が与えられるので、 $\displaystyle \sum_{i=1}^N \sum_{j=i}^N f(A_i+A_j)$ を求めてください。

• $1\le N\le 2\times 10^5$
• $1\le A_i\le 10^7$
• 入力は全て整数

$A$ のbincount（各 $A_i$ に対して $C[A_i]+=1$ した配列 $C$）を2つ畳み込めば、 合計値 $x$ ごとに、$x=A_i+A_j$ となる $(i,j)$ ペアの個数が分かるが、、、
その解法は防ぐためか、制約が $A_i \le 10^7$ となっていて時間的に厳しい。
（C++ ならあるいは？と思ったが、C++のatcoderライブラリのconvolution_ll()は 畳み込みサイズに上限があり、上手くいかなかった）

下のbitから、偶奇でグループを分割しながら調べていく。
例えば偶数+奇数は必ず奇数なのでそのbitで確定するためまとめて計算でき、 偶数+偶数など、より上位の桁まで0が続くものは再帰で調べていく、という感じ。

以下の2つの再帰関数を実装すれば良い。

• $f(A,carry)$
  • 1つの数列内の要素同士を組み合わせて作られるペアについて調べる。
  • carryは、下の桁からの繰り上がりの有無(0/1)を示す。
• $g(A,B)$
  • 2つの数列から1つずつ取り出して作られるペアについて調べる。

$f$ は、まず $A$ を偶奇で2つの集合（$Even,Odd$）に分ける。

• $carry=0$ のとき
  • 偶+奇ペアがこのbitで確定する。
    • （Evenの総和×Oddの個数）＋（Oddの総和×Evenの個数）が答えに加算される。
  • Even同士のペアは、各要素を1bit右シフトし、$f(Even, 0)$
  • Odd同士のペアは、各要素を1bit右シフトし、$f(Odd, 1)$
• $carry=1$ のとき
  • Even同士ペア、Odd同士ペアが確定する。
    • Evenなら、Evenの総和×（Evenの個数+1）＋（Even同士でできるペアの個数）が答えに寄与する。
    • 2項目は、繰り上がりのためにペア毎に1ずつ加算される分。
  • 偶+奇ペアは、各要素を1bit右シフトし、$g(Even,Odd)$

$g$ は、$A,B$ それぞれで偶奇の集合（$A_E,A_O,B_E,B_O$）に分ける。
性質上、必ず $carry=1$ である。

• 偶数同士（$A_E,B_E$）・奇数同士（$A_O,B_O$）のペアが確定する。
  • 偶数同士は（$A_E$ の総和×$B_E$ の個数）＋（$B_E$ の総和×$A_E$ の個数）＋（$A_E$ の個数×$B_E$ の個数）
    • 最後の項は、繰り上がりのためペア毎に1ずつ加算される分。
  • 奇数同士も同様に求められる
• 偶奇ペアは、各要素を1bit右シフトしてそれぞれ再帰。$g(A_E,B_O)$、$g(A_O,B_E)$

実際は、$A$ には同じ要素が含まれうるため｛値: 個数｝の辞書型で管理するとよい。

各要素が再帰を繰り返す回数は $\log{A_i}$ に抑えられ、1回の再帰関数では配列長に比例するので、計算量は $O(N \log{A_{\max}})$ となる。

Python3

• 公式Editorial

たとえば2進数で下4桁が $A_i=...1101$ である $A_i$ に対しては、 下4桁が $...0011$ である $A_j$ が、 「$A_i+A_j$ が2で少なくとも4回割り切れるペア」となる。
（つまり、bit反転して1足した値）

$d_k:=$「$A_i+A_j$ が2で少なくとも $k$ 回割り切れるペアの、$A_i+A_j$ の総和」とする。
すると、$d_k-d_{k+1}$ が、「2でちょうど $k$ 回割り切れるペアの $A_i+A_j$ の総和」となる。
そのようなペアの答えへの寄与は、$\dfrac{d_k-d_{k+1}}{2^k}$ となる。

よって、$d_0,d_1,...$ が求められればよい。

$d_k$ は、$A_i$ の下 $k$ 桁ごとに個数と総和をカウントすることで求められる。

G - Abs Sum

• 長さ $N$ の整数列 $A=(A_1,A_2,\ldots, A_N)$ 、 $B=(B_1,B_2,\ldots,B_N)$ と長さ $K$ の整数列 $X=(X_1,X_2,\ldots,X_K)$ 、 $Y=(Y_1,Y_2,\ldots,Y_K)$ が与えられます。
• $k=1,2,\ldots,K$ に対して $\displaystyle \sum_{i=1}^{X_k} \sum_{j=1}^{Y_k} \left|A_i-B_j \right|$ を求めてください。

• $1\le N\le 10^5$
• $0\le A_i,B_j\le 2\times 10^8$
• $1\le K\le 10^4$
• $1\le X_k,Y_k\le N$
• 入力は全て整数

$K$ はクエリの個数、$X_i,Y_i$ は1つのクエリと見なせる。

まず、クエリが1つ（たとえば $A$ と $B$ 全体に対して求める）の場合を考える。
各 $B_j$ について、$B_j$ 以下の $A_i$ とは $B_j-A_i$ で答えに寄与し、$B_j$ より大きい $A_i$ とは $A_i-B_j$ で寄与する。

よって、$B_j$ より小さい/大きい $A$ の要素の総和と個数が分かれば良い。
$A$ をソート・累積和を前計算しておいて、$B_j$ ごとに $A$ を二分探索することで、$O(N\log{N})$ で求められる。

$A_1～A_{i}$ のみの「ソートされ、二分探索と累積和の取得ができる状態」が、$i$ を変化させても保持できれば、 $B_j$ を1つ加える/削除することの寄与を差分計算できる。
（同様に、$B$ に対しても同じものを用意すれば、$A_i$ を1つ加える/削除する操作も可能）

これは、座標圧縮した上で、Fenwick Tree 2本で総和と個数をそれぞれ管理すると可能である。
$A,B$ いずれも、右端を1つ伸ばして/縮めて、差分計算をする、という操作を $O(\log{N})$ でおこなえる。

クエリ平方分割（Mo's Algorhithm）をする。

• Mo's algorithm - ei1333の日記

バケットサイズを $S$ とし、$\left \lfloor \dfrac{X_i}{S} \right \rfloor$ が同じクエリを1つのバケットにまとめる。
バケットごとに答えを求める。1つのバケット内では、$Y_i$ が小さいものから求めるようにする。

        着目中のバケット
|        |----S----|       |
          <------->             A の右端は Xi によってクエリ毎にこの範囲を行ったり来たりする
 ------------------------->     B の右端はバケット内の処理で一様に右に移動する

$A$ の右端が動く回数は、クエリ毎に最大 $S$ 動くので、全体で $O(KS)$ 回。
$B$ の右端が動く回数は、バケットが $\dfrac{N}{S}$ 個あり、それぞれで $O(N)$ かかるので $O(\dfrac{N^2}{S})$ 回。

この2つを大体同じにすると最悪の場合の計算量が抑えられる。$S=\dfrac{N}{\sqrt{K}}$ くらいにすればよい。
（実際は、1回のバケット毎にFenwick Treeを用意したりなどの付随処理がかかるため、 $S$ は少し大きめに、バケット数が少なくなるように調整した方が高速になりやすい）

右端が1つ動く毎に $O(\log{N})$ の計算量がかかるので、全体で $O(N\sqrt{K}\log{N})$ となる。

制限時間が5秒とはいえ、遅い言語だとちょっと下手な書き方をしてしまうと厳しい。
バケットサイズの調整や、次のバケットの計算に移るときに前のバケットの結果を再利用できるところはする、 bitshiftして総和と個数を1つの整数にまとめてFenwickTreeを $A,B$ につき各1本で済ます、などの高速化が効く場合がある。

Python3