问题

Problem Statement

给出一个长度为 N 的序列 A ，其由 1 到 \textbf{10} 之间的整数组成。

如果满足以下条件，那么三元组 (l,r) （ 1\leq l \leq r\leq N ）是“好数对”：

• 序列 (A_l,A_{l+1},\ldots,A_r) 包含长度为 3 的等差子序列。换言之，存在整数对 (i,j,k) （ l \leq i < j < k\leq r ），使得 A_j - A_i = A_k - A_j 。

求“好数对”的个数。

Constraints

• 3 \leq N \leq 10^5
• 1\leq A_i \leq 10
• 所有输入均为整数。

Input

输入格式如下：

• N
• A_1 \ldots A_N

Output

输出“好数对”的个数。

解法

思路

利用双指针的思想，将 l 和 r 分别看作左右指针。

考虑到 A_i 的范围很小，可以直接使用队列存储三元组的第一个数；再利用前缀和存储每个数的总数，用于处理三元组的第二个数。

然后只需要从头开始遍历 r ，找三元组的第三个数。每遍历到一个数，就将 r 的值推入到这个数对应的队列内，紧接着在符合条件的非空队列中提取首值，然后计算对应的中间值，利用之前保存的前缀和判断从 l+1 到 r-1 的区间内是否存在中间值。

如果存在，说明目前从 l 到 r 的区间中存在符合题意的三元组，再根据 A_r 到 A_n 的距离确定当前 l 对应的三元组总数量，并将其添加到答案中；然后弹出 l 所在队列的首值（弹出的值应该正好为l），自增 l 后重复上述步骤，直到中间值不存在或对应队列被清空。

遍历完 r 之后即可输出答案，一定记得开 long long ，否则会溢出。时间复杂度为 O(n)。

实现

按如下方式开前缀和数组和队列：

int r[11][100010];
queue<int> st[11];

以下为实现：

--snip--
void solve() {
  int ptr = 1; // ptr 表示 l 的位置
  for (int i = 1; i <= n; ++i) {
    st[arr[i]].push(i); // 在对应队列推入当前的 r 值
    for (int j = 1; j <= 10; ++j) {
      // 遍历队列，若队列为空或中间值不为整数则跳过
      if (st[j].empty() || (j + arr[i]) % 2 != 0) continue;
      int p = st[j].front(); // 指定三元组第一个数的位置
      // 通过前缀和判定中间值是否存在
      while ((j + arr[i]) % 2 == 0 &&
             r[(j + arr[i]) / 2][i - 1] - r[(j + arr[i]) / 2][p] > 0) {
        ans += n - i + 1; // 从 (l, r) 到 (l, n) 都包括此三元组
        // 弹出左指针对应的元素，并令其右移
        st[arr[ptr]].pop();
        ptr++;
        if (st[j].empty()) break; // 对应队列为空时跳出 while 循环
        p = st[j].front();
      }
    }
  }
}

int main(){
  cin >> n;
  for (int i = 1; i <= n; ++i) {
    cin >> arr[i];
    // 根据输入值处理前缀和，非输入值的前缀和维持不变
    for (int j = 1; j <= 10; ++j) {
      if (j == arr[i])
        r[arr[i]][i] = r[arr[i]][i - 1] + 1;
      else
        r[j][i] = r[j][i - 1];
    }
  }
  solve();
  cout << ans << endl;
  return 0;
}