B - Arithmetic Progression Subsequence

Problem Statement

You are given a sequence $A$ of length $N$ consisting of integers between $1$ and $\textbf{10}$, inclusive.

A pair of integers $(l,r)$ satisfying $1\leq l \leq r\leq N$ is called a good pair if it satisfies the following condition:

• The sequence $(A_l,A_{l+1},\ldots,A_r)$ contains a (possibly non-contiguous) arithmetic subsequence of length $3$. More precisely, there is a triple of integers $(i,j,k)$ with $l \leq i < j < k\leq r$ such that $A_j - A_i = A_k - A_j$.

Find the number of good pairs.

Constraints

• $3 \leq N \leq 10^5$
• $1\leq A_i \leq 10$
• All input numbers are integers.

Input

The input is given from Standard Input in the following format:

$N$
$A_1$ $\ldots$ $A_N$

Output

Print the answer.

Sample Input 1

5
5 3 4 1 5

Sample Output 1

3

There are three good pairs: $(l,r)=(1,4),(1,5),(2,5)$.

For example, the sequence $(A_1,A_2,A_3,A_4)$ contains an arithmetic subsequence of length $3$, which is $(5,3,1)$, so $(1,4)$ is a good pair.

Sample Input 2

3
1 2 1

Sample Output 2

0

There may be cases where no good pairs exist.

Sample Input 3

9
10 10 1 3 3 7 2 2 5

Sample Output 3

3

解题思路

  唉昨晚爆 0 了。先提一下官方给出的做法，太 trick 了我是不可能想到的。

  事实上不满足条件的数对 $(l,r)$ 的长度（即 $r-l+1$）一定不超过 $20$。因为 $1 \leq a_i \leq 10$，由鸽巢原理知道如果某个 $(l,r)$ 的长度超过 $20$，那么必然存在某个值出现了至少 $3$ 次，此时把这个值的 $3$ 个下标作为 $(i,j,k)$ 那么一定有 $a_j - a_i = a_k - a_j = 0$。即如果 $(l,r)$ 的长度超过 $20$，那么这些数对一定是满足条件的。因此我们只需枚举所有长度不超过 $20$ 的数对，并判断是否存在满足条件的 $(i,j,k)$ 即可，时间复杂度大概是 $O(10^2 \cdot n)$。

  下面给出更容易想到的做法。其实看到这个数据范围就知道做法肯定与值域有关。

  枚举 $(l,r)$ 中的 $r$，然后判断在 $[1,r]$ 中是否存在满足条件的 $(i,j,k)$。暴力的思路就是把所有满足条件的 $(i,j,k)$ 找出来，并找到当中 $i$ 的最大值，记作 $t$。那么以 $r$ 为右端点且满足条件的数对数量就是 $t$。暴力求肯定会超时，而且也不容易想到如何用值域来优化。

  因此想到能否预处理出每个前缀中满足条件的 $(i,j,k)$ 的 $i$ 的最大值。把所有满足条件的 $(i,j,k)$ 按照 $j$ 分成 $n$ 类，对于每个 $j$，我们不枚举下标，而是枚举值 $1 \leq a_i \leq 10$，$1 \leq a_k \leq 10$。如果满足 $a_j - a_i = a_k - a_j$，那么我们应该找到在 $[1, j-1]$ 中是否存在值 $a_i$，如果存在则选择最大下标。同理在 $[j+1, n]$ 中是否存在值 $a_k$，如果存在则选择最小下标（贪心，显然 $k$ 越小则能被更多的 $(l,r)$ 包含）。如果都存在，我们令 $f(R_{j+1,a_k}) = \min\{ f(R_{j+1,a_k}), \, L_{j-1, a_i} \}$。其中 $L_{j-1, a_i}$ 是 $[1, j-1]$ 中值 $a_i$ 出现的最大下标，$R_{j+1, a_k}$ 是 $[j+1, n]$ 中值 $a_k$ 出现的最小下标，这两个都可以通过 $O(10 \cdot n)$ 的复杂度预处理出来。$f(x)$ 则表示所有满足条件且 $k=x$ 的 $(i,j,k)$ 中，$i$ 的最小值。

  上面预处理的时间复杂度是 $O(10^2 \cdot n)$，当然还可以优化到 $O(10 \cdot n)$，即只枚举 $a_i$，$a_k$ 通过 $a_k = 2 \cdot a_i - a_j$ 得到，只需判断是否满足 $1 \leq a_k \leq 10$ 即可。

  最后由于求的是前缀中 $i$ 的最小值，只需从左到右遍历令 $f(i) = \min\{ f(i), f(i-1) \}$ 即可。

  AC 代码如下，时间复杂度为 $O(10 \cdot n)$：

#include <bits/stdc++.h>
using namespace std;

typedef long long LL;

const int N = 1e5 + 10;

int a[N];
int l[N][11], r[N][11];
int f[N];

int main() {
    int n;
    scanf("%d", &n);
    for (int i = 1; i <= n; i++) {
        scanf("%d", a + i);
    }
    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= 10; j++) {
            l[i][j] = l[i - 1][j];
            if (a[i] == j) l[i][j] = i;
        }
    }
    for (int i = n; i; i--) {
        for (int j = 1; j <= 10; j++) {
            r[i][j] = r[i + 1][j];
            if (a[i] == j) r[i][j] = i;
        }
    }
    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= 10; j++) {
            int k = 2 * a[i] - j;
            if (k < 1 || k > 10) continue;
            if (l[i - 1][j] && r[i + 1][k]) f[r[i + 1][k]] = max(f[r[i + 1][k]], l[i - 1][j]);
        }
    }
    for (int i = 1; i <= n; i++) {
        f[i] = max(f[i], f[i - 1]);
    }
    LL ret = 0;
    for (int i = 1; i <= n; i++) {
        ret += f[i];
    }
    printf("%lld", ret);
    
    return 0;
}

参考资料

  Editorial - AtCoder Regular Contest 170：https://atcoder.jp/contests/arc170/editorial/9148

  Submission #49545473 - AtCoder Regular Contest 170：https://atcoder.jp/contests/arc170/submissions/49545473