【ABC409】C問題 - Equilateral Triangle 考察から実装(c++)まで

AtCoderBeginnerContest409の解説＆感想です。
コンテストリンク

問題一覧

• 【ABC409】A問題 - Conflict　考察から実装(c++)まで
• 【ABC409】B問題 - Citation　考察から実装(c++)まで
• 【ABC409】C問題 - Equilateral Triangle　考察から実装(c++)まで <- イマココ
• 【ABC409】D問題 - String Rotation　考察から実装(c++)まで
• 【ABC409】E問題 - Pair Annihilation　考察から実装(c++)まで
• 【ABC409】F問題 - Connecting Points　考察から実装(c++)まで

C問題 - Equilateral Triangle

問題リンク

問題概要

長さ$L$の円周上に、$N$個の点$1,2,...,N$がある。点$i+1$は点$i$から時計回りに$d_i$進んだ位置にある。
$a \lt b \lt c$であって、$a,b,c$を直線で結ぶと正三角形になるような整数の組$(a,b,c)$の個数を求めよ。

制約

• $ 1 \le N,L \le 3 \times 10^5 $
• $ 0 \le d_i \le L$

考察

わりに素直に考えればいい。
まず、$L$が三等分できない時は正三角形が作れないので、最初にはじいておく。
あとは一辺$\frac{L}{3}$の正三角形がいくつあるか？が分かればいい。

円周上の適当な場所$P$を一つ固定して、その場所にある点の個数を$cnt(P)$とすると、$P$から時計回りにできる正三角形は

• $cnt(P) \times cnt((P+\frac{L}{3})\%N) \times cnt((P+2\frac{L}{3})\%N)$

で数えられる。これは$cnt$を配列なりmapなりで取っておけば$O(1)$で計算できるので、円周上の$L$個の場所すべてに対して上式を計算して、$3$で割ればいい。（同じ正三角形を$3$回数えてしまうため）
計算量は、$cnt$の作成と全ての場所をなめるループがそれぞれあるので$O(N+L)$。

実装

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using vll = vector<ll>;

int main(void){
    //入力
    ll N,L;
    cin>>N>>L;
    
    vll D(N-1);
    for(int i=0;i<N-1;i++) cin>>D[i];
    
    //そもそも正三角形が作れなければアウト
    if(L%3 != 0){
        cout<<0<<endl;
        return 0;
    }
    
    //点1の場所を0として、頻度マップを作る
    vll cnt(L,0);
    ll nowP = 0;
    cnt[nowP]++;
    for(int i=0;i<N-1;i++){
        nowP = (nowP + D[i]) % L;
        cnt[nowP]++;
    }
    
    //円周上の各点に対して、そこから時計回りにできる正三角形の個数を数えていく
    ll ans = 0;
    for(int p = 0;p<L;p++) ans += cnt[p]*cnt[(p+L/3)%L]*cnt[(p+2*L/3)%L];
    cout<<ans/3<<endl;
}