公式解説では主客転倒 + 余事象を使った解き方が紹介されていましたが、少し違うアプローチ（分類するならDP?）で解いたのでそれをまとめておきたいと思います。
問題
https://atcoder.jp/contests/abc371/tasks/abc371_e
問題概要
長さ Nの整数列A=(A_{1}, \dots, A_{N})に対し、f(l, r) \ (1\leq l \leq r \leq N)を、「集合\{A_{l}, A_{l+1}, \dots, A_{r-1}, A_{r}\} の要素数」と定める。このとき
\sum_{i=1}^{N}\sum_{j=i}^{N}f(i, j)を求めよ。
制約
- 1\leq N \leq 2\times 10^{5}
- 1\leq A_{i} \leq N
解法
以下ではAの添え字を一つずらして0-indexed で考えます。
方針
g(i):=\sum_{j=i}^{N}f(i, j) と書くことにします。最終的に求めるものはg(0) + \cdots + g(N-1)です。ここで
-
g(0) の計算はO(N) で可能。
- よって、各1\leq i \leq N-1について、g(i-1)からg(i) を高速に求められれば良い。
-
g(i)とg(i-1)の差は、A_{i-1}と、A_{k}=A_{i-1} \ (k\geq i)なるkの情報が分かれば良い。
ということに着目しました。
詳細
まず\text{appear}[i][x]を、「i\leq k \leq N-1かつA_{k} = xとなるk」を昇順に並べたリストとします。
例えば、N=3, A=(1, 2, 2) (問題の入力例1)の場合は次の表のようになります：
| i \setminus x |
1 |
2 |
| 0 |
[1] |
[2, 3] |
| 1 |
[] |
[2, 3] |
| 2 |
[] |
[3] |
このとき、1\leq i\leq N-1について次のことが成り立ちます：
-
\text{appear}[i][A_{i-1}] が空のとき、g(i) = g(i-1) - (N-i+1).
-
\text{appear}[i][A_{i-1}] が空でないとき、g(i) = g(i-1) - (\text{appear}[i][A_{i-1}][0]-i+1).
証明
まず1について、\text{appear}[i][A_{i-1}] が空であるとします。この場合、j\geq iであればA_{j} \ne A_{i-1} が成り立ちます。よって、各j = i, i+1, \dots, N-1について、
\{A_{i-1}, A_{i}, A_{i+1}, \dots, A_{j}\} = \{A_{i-1}\} \sqcup \{A_{i}, A_{i+1}, \dots, A_{j}\}
が成り立ちます。よって、
\begin{aligned}
g(i) &= \sum_{j=i}^{N-1}\left(\text{集合}\{A_{i}, A_{i+1}, \dots A_{j}\}\text{の要素数}\right) \\
&= \sum_{j=i}^{N-1}\left(\text{集合}\{A_{i-1}, A_{i}, A_{i+1}, \dots A_{j}\}\text{の要素数} - 1\right) \\
&= \sum_{j=i}^{N-1}\left(\text{集合}\{A_{i-1}, A_{i}, A_{i+1}, \dots A_{j}\}\text{の要素数}\right) - (N-i) \\
&= \sum_{j=i-1}^{N-1}\left(\text{集合}\{A_{i-1}, A_{i}, A_{i+1}, \dots A_{j}\}\text{の要素数}\right) - \left(\text{集合}\{A_{i-1}\}\text{の要素数}\right) - (N-i) \\
&= g(i-1) - (N-i+1)
\end{aligned}
が成り立ちます。
2についても考え方は同様です。\text{appear}[i][A_{i-1}] が空でないとします。このとき、\text{appear}[i][A_{i-1}][0]がA_{k}=A_{i-1}かつk\geq iを満たす最小のkであることから、K:=\text{appear}[i][A_{i-1}][0] とすると、
\begin{aligned}
i\leq j < K &\Rightarrow \{A_{i-1}, A_{i}, A_{i+1}, \dots, A_{j}\} = \{A_{i-1}\} \sqcup \{A_{i}, A_{i+1}, \dots, A_{j}\}, \\
K \leq j \leq N-1 &\Rightarrow \{A_{i-1}, A_{i}, A_{i+1}, \dots, A_{j}\} = \{A_{i}, A_{i+1}, \dots, A_{j}\}
\end{aligned}
が成り立ちます。従って、
\begin{aligned}
g(i) &= \sum_{j=i}^{N-1}\left(\text{集合}\{A_{i}, A_{i+1}, \dots A_{j}\}\text{の要素数}\right) \\
&= \sum_{j=i}^{K-1}\left(\text{集合}\{A_{i}, A_{i+1}, \dots A_{j}\}\text{の要素数}\right) \\
& \ \ \ \ + \sum_{j=K}^{N-1}\left(\text{集合}\{A_{i}, A_{i+1}, \dots A_{j}\}\text{の要素数}\right)\\
&= \sum_{j=i}^{K-1}\left(\text{集合}\{A_{i-1}, A_{i}, A_{i+1}, \dots A_{j}\}\text{の要素数} - 1\right) \\
& \ \ \ \ + \sum_{j=i}^{K-1}\left(\text{集合}\{A_{i-1}, A_{i}, A_{i+1}, \dots A_{j}\}\text{の要素数}\right) \\
&= \sum_{j=i}^{N-1}\left(\text{集合}\{A_{i-1}, A_{i}, A_{i+1}, \dots A_{j}\}\text{の要素数}\right) - (K-i) \\
&= \sum_{j=i-1}^{N-1}\left(\text{集合}\{A_{i-1}, A_{i}, A_{i+1}, \dots A_{j}\}\text{の要素数}\right) - \left(\text{集合}\{A_{i-1}\}\text{の要素数}\right) - (K-i) \\
&= g(i-1) - (K-i+1) \\
&= g(i-1) - (\text{appear}[i][A_{i-1}][0]-i+1)
\end{aligned}
が成り立ちます。
よってg(i-1), \text{appear}[i][A_{i-1}] が分かっていれば、g(i)はO(1)で求まります。
今、\text{appear}[0]はO(N)で求めることができます。また、
\text{appear}[i][x] = \begin{cases}
\text{appear}[i-1][x], &\text{if } x \ne A_{i-1} \\
\text{appear}[i-1][x] \text{から先頭の要素を削除したもの}, &\text{if } x = A_{i-1}
\end{cases}
なので、in-place に更新することで\text{appear}[i-1]から\text{appear}[i-1]はO(1)で求まります。
以上から、g(0), g(1), \dots g(N-1) がO(N)で求まり、最終的な答えである
も
O(N)で求まります。
実装
Python での実装（コンテスト後にこの記事に合わせて少し書き換えたもの）です。
https://atcoder.jp/contests/abc371/submissions/57850582
(ほとんど同じではありますが、コンテスト中の実装よりはよくなっていると思います。)
Discussion