ABC409 C – Equilateral Triangle

ギリギリになって、累積和で点の座標を求めるところまで思いついたのに、正三角形を求める方法が思いつきませんでした。悔しい、、

今回も解説を読んでいきます。

C - Equilateral Triangle

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

atcoder.jp

問題

円周上の点1,2,…,Nが配置されていて、点i+1は点iから時計回りに円周上をdi進んだ位置にあります。

その中で3点が全て異なる位置で、その３点で正三角形であるものの個数を求めます。

解法

入力例1で考えます。

まずは各点の座標xを求めます。

例えばi=2の時、x[2]はi=1を0として4進んだ位置にいるためx[2]=4です。

i=3であれば、7進んだ位置になりますが、円なので座標としては1になります。

この要領でそれぞれ求めると以下の表のようになります。

上記はx[i] = (d[1]+ d[2] + ... + d[i-1]) mod Lで求まります。

xがわかったらあとは3点が正三角形になるものを求めます。

これは円周上の３点（a,b,c）が等間隔に並んでいるものになります。

xa < xb < xcの時以下の条件式で表すことができます。

xb = xa + L/3 かつ xc = xb + L/3

より一般にすると

xa,  xb, xc = k, k+L/3, 2L/3 ## Lが6であれば 0, 2, 4 や 1, 3, 5 などになります

ここでxの頻度配列をcntとします。

cntは入力例1では以下のようになります。

例えば上記で正三角形になるのは0,2,4の点の時なので、1*1*2=2で2つの正三角形ができます。

また1,3,5の時は3と5に点がないので正三角形はできません。

これをxが0からL//3まで繰り返せば良いです。

以下のコードでACすることができました。

def i_map():
    return map(int, input().split())


def i_list():
    return list(map(int, input().split()))


n, l = i_map()
d = i_list()

if l % 3 != 0:
    print(0)
    exit()

x = 0
cnt = [0] * l
for i in range(n):
    if i != 0:
        x = (x + d[i - 1]) % l
    cnt[x] += 1

ans = 0
for i in range(l // 3):
    ans += cnt[i] * cnt[i + l // 3] * cnt[i + 2 * l // 3]

print(ans)