あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

Indeedなう（予選B） D - 高橋くんと数列

競プロ

問題原文

解法

とりうる $(l, r)$ の組み合わせは $\frac{N × (N + 1)}{2}$ 通り。（ $l = r$ でも良いので）

各 $c$ について、条件を満たすものを数え上げようとすると、結局 $1 ~ N$ のループを各 $c$ について回したくなるし、なにより面倒なので 条件を満たさないものを数えて引く ことを考える。

条件を満たさないものは以下のように計算できる。
f:id:AT274:20200525225520j:plain

実装

先頭と末尾に番兵を置いておくと楽です。

from collections import defaultdict
N, C = map(int, input().split())
A = list(map(int, input().split()))

ALL = N * (N + 1) // 2
D = defaultdict(list)
for i, a in enumerate(A):
    D[a].append(i)
for c in range(1, C + 1):
    D[c] = [-1] + D[c] + [N]

for c in range(1, C + 1):
    ans = 0
    for i in range(len(D[c]) - 1):
        ans += (D[c][i + 1] - D[c][i]) * (D[c][i + 1] - D[c][i] - 1) // 2
    print(ALL - ans)

感想

その数を１つでも含むような〜 → その数を１つも含まない〜 の言い換えは受験とかでもよくみますね。