Official

C - Equilateral Triangle Editorial by toam


はじめに,各点 \(i=1,2,\ldots,N\) に対して点 \(1\) と点 \(i\) の相対的な位置関係を求めます.点 \(1\) の場所を座標 \(0\) として,そこから時計回りに \(d\ (0\leq d\lt L)\) 進んだ位置にある場所を座標 \(d\) とします.

\(i\) は点 \(1\) から時計方向に \(d_1+d_2+\dots+d_{i-1}\) だけ進んだ距離にあるため,その座標 \(x_i\)\(x_i=d_1+d_2+\dots+d_{i-1}\ (\mathrm{mod}\ L)\) です.累積和を用いることで,\(x_1,x_2,\ldots,x_N\)\(O(N)\) で求めることができます.

\(3\)\(a,b,c\) が正三角形をなすということは,この \(3\) 点が円周上に等間隔に並んでいるということです.たとえば \(x_a<x_b<x_c\) のとき,この条件を数式で表すと

  • \(x_b=x_a+L/3\) かつ \(x_c=x_b+L/3\)

です.より一般に,\(3\) 点が等間隔に並ぶ条件は次のように表せます.

  • ある整数 \(k\ (0\leq k\lt L/3\) が存在して,\(\lbrace x_a,x_b,x_c\rbrace=\lbrace k,k+L/3,k+2L/3\rbrace\)

したがって,\(x_1,x_2,\dots,x_N\) の頻度配列を \(\mathrm{cnt}\) とすれば,求める答えは \(\displaystyle \sum_{k=0}^{L/3-1} \mathrm{cnt}[k]\times \mathrm{cnt}[k+L/3]\times \mathrm{cnt}[k+2L/3]\) です.

\(L\)\(3\) で割り切れないときは答えは \(0\) になります.

実装例(Python)

N, L = map(int, input().split())
if L % 3 != 0:
    print(0)
    exit()
d = list(map(int, input().split()))
x = 0
cnt = [0] * L
for i in range(N):
    if i != 0:
        x += d[i - 1]
    x %= 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)

posted:
last update: