【AtCoder】ABC 409 C - Equilateral Triangle | 緑コーダーが解くAtCoder

atcoder.jp

配点: 300 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 421 / NoviSteps: 2Q

問題概要
制約
考察
実装例
コメント

問題概要

円周が $L$ の円があり、この円周上に $N$ 個の点が配置されている。 $i=1,2,\dots,N-1$ に対し、点 $i+1$ は点 $i$ から時計回りに円周上を $d_i$ 進んだ位置にある。

整数の組 $(a,b,c) \: (1\leq a\lt b\lt c\leq N)$ であって、以下の2つをともに満たすものの個数を求めよ。

3点 $a,b,c$ はすべて異なる位置にある。
3点 $a,b,c$ を頂点とする三角形は正三角形である。

制約

$3\leq L,N\leq 3\times 10^5$
$0\leq d_i\lt L$
入力は全て整数

考察

大前提として、円周上の3点を選んで正三角形を作るためには $L$ を3等分する必要がある。つまり、 $L$ は $3$ の倍数でなければならない。

その上で、点 $1$ を原点として円周上に座標軸を取り、その座標系での $N$ 個の点の座標 $x_i$ を以下のように累積和を利用して求める。

$\begin{align*} x_i = (x_{i-1} + d_{i-1}) \mod L \quad (i = 2, 3, \dots, N) \end{align*}$

これらの座標の中から、座標が $x_a, x_b, x_c$ であるような3点を選ぶ $(x_a \lt x_b \lt x_c)$ 。

これら3点の関係は以下のようになる。

$x_b = x_a + \frac{L}{3}$
$x_c = x_a + \frac{2L}{3}$

このとき、各座標ごとの点の度数分布表をmapで管理すれば(これを $\mathrm{cnt}$ とする)、 $x_a$ を固定したときに作ることができる正三角形の個数は

$\begin{align*} \mathrm{cnt}_{x_a} \times \mathrm{cnt}_{x_a + \frac{L}{3}} \times \mathrm{cnt}_{x_a + \frac{2L}{3}} \end{align*}$

と計算できる。したがって、 $x_a$ の位置を $0$ から $\frac{L}{3}$ まですべて試して、その総和を取ったものが答えとなる。

実装時はオーバーフローに注意(1敗)。

実装例

#include <bits/stdc++.h>

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

using namespace std;
using ll = long long;

// ======================================== //

int main()
{
    ll N, L;
    cin >> N >> L;
    vector<int> d(N - 1);
    rep(i, 0, N - 1) cin >> d[i];

    if (L % 3 != 0)
    {
        cout << 0 << endl;
        return 0;
    }

    vector<ll> pos(N, 0);
    map<ll, ll> cnt;
    cnt[0]++;
    rep(i, 1, N)
    {
        pos[i] = (pos[i - 1] + d[i - 1]) % L;
        cnt[pos[i]]++;
    }

    ll ans = 0;
    ll distance = L / 3;
    rep(i, 0, distance)
    {
        ans += cnt[i] * cnt[i + distance] * cnt[i + 2 * distance];
    }

    cout << ans << endl;

    return 0;
}