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乗したものを足したものが答えになる。