ARC 102 | C - Triangular Relationship

問題

https://atcoder.jp/contests/arc102/tasks/arc102_a

解法

問題文より以下である。

$$a + b \equiv b + c \equiv c + a \equiv 0 \ (\bmod \ K)$$

ここで$a + b \equiv b + c$から$b$を引くと

$$a \equiv c \ (\bmod \ K)$$

対称性より

$$a \equiv b \equiv c \ (\bmod \ K)$$

上記を利用し、$a + b \equiv 0$より

$$a + b \equiv 2a \equiv 0 \ (\bmod \ K)$$

ここで$K$が奇数のとき、$2$と$K$は互いに素なので$a$は$K$で割り切れる。

$K$が偶数の場合、$K = 2K'$と置くと

$$2a \equiv 0 \ (\bmod \ K)$$

$$2a \equiv 0 \ (\bmod \ 2K')$$

$$a \equiv 0 \ (\bmod \ K')$$

対称性より

$$a \equiv 0 \ (\bmod \ K')$$

$$b \equiv 0 \ (\bmod \ K')$$

$$c \equiv 0 \ (\bmod \ K')$$

よって

$$a \equiv b \equiv c \equiv 0 \ (\bmod \ K)$$

または$K = 2K'$なので

$$a \equiv b \equiv c \equiv K' \ (\bmod \ K)$$

コード

#include <bits/stdc++.h>

#include <atcoder/all>

using namespace std;
using namespace atcoder;
using ll = long long;
using ld = long double;
using uint = unsigned int;
using ull = unsigned long long;
const int MOD = 1e9 + 7;

int main() {
  ll n, k;
  cin >> n >> k;
  vector<ll> v(k);
  for (int i = 1; i <= n; i++) {
    v[i % k]++;
  }
  if (k % 2 == 0) {
    cout << v[0] * v[0] * v[0] + v[k / 2] * v[k / 2] * v[k / 2] << endl;
  } else {
    cout << v[0] * v[0] * v[0] << endl;
  }
}

参考

https://atcoder.jp/contests/arc102/tasks/arc102_a/editorial