phyllo’s algorithm note

レッドコーダーへの道のりは遠い。休んでる場合じゃない!

ARC102 C. Triangular Relationship

問題

整数N, Kが与えられ、N以下の正の整数の組(a, b, c)を考える。
a+b, b+c, c+aがKの倍数であるような(順番を考慮する)組の個数を答えよ。

制約

  • 1 <= N, K <= 2*10^5

解法

2つの整数a, bをループで回してcの候補数を列挙するような方法だとO(N^2)かかってしまい間に合わない。


まず、「Kの倍数」をそのまま扱うのは大変なので、mod Kで考える。
そうすると、「Kの倍数である」<=>「mod Kで0になる」になる。


そこで、mod Kの世界で考え、(a mod K, b mod K, c mod K)の組み合わせを考えると、2つとって足したら0になる必要があるため、実は以下の場合しかない。

  • 0を足す場合:(0, 0, 0)
  • 0でないものを足す場合:(K/2, K/2, K/2) (ただし、K/2が割り切れること)

したがって、Kが奇数の時は(0, 0, 0)のパターンだけで、「1~Nまででmod Kが0になる個数」が各a,b,cの位置にあり得るので、個数を3乗したものが答え。
Kが偶数の時は、(0, 0, 0)のパターンと(K/2, K/2, K/2)のパターンがあり、同様にそれぞれ個数を3乗したものを足したものが答えになる。