[ABC318E] Sandwiches 题解

题意

给定一个长度为 \(N\) 的正整数列 \(A = \left(A_1, A_2, \cdots,A_N\right)\)，求满足以下条件的正整数三元组 \(\left(i, j, k\right)\) 的数量：

• \(1 \le i < j < k \le N\)；
• \(A_i = A_k\)；
• \(A_i \neq A_j\)。

数据范围：

• \(3 \le N \le 3 \times 10^5\)；
• \(1 \le A_i \le N\)。

题解

容斥做法。

在计数过程中先不考虑第三个限制 \(A_i \neq A_j\)，最后再去除不合法情况。

现在要处理的限制只有 \(1 \le i < j < k \le N\) 和 \(A_i = A_k\)，发现对 \(j\) 限制很少，故考虑枚举 \(k\)，统计合法的 \(i\)，然后计算出 \(j\) 的数量。设存在 \(i,k\) 使得 \(A_i = A_k\)，那么所有 \(j \in \left(i, k\right)\) 均合法，数量为 \(k - i - 1\)。可以快速计算。

设 \(f_x = \sum\limits_{i} 1\left[A_i = x\right], g_x = \sum\limits_{i} i\left[A_i = x\right]\)，从 \(1\) 到 \(N\) 枚举 \(k\) 的值并计算前缀的 \(f\) 值和 \(g\) 值，对于 \(A_k = x\)，那么这个 \(k\) 对答案的贡献即

\[\left(k - 1\right) \times f_x - g_x \]

这一部分可以 \(\mathcal{O}(N)\) 完成计算，下面考虑去除不合法情况。

发现不合法的三元组 \(\left(i, j, k\right)\) 一定满足 \(A_i = A_j = A_k = x\)，故枚举 \(x\) 值，计算不合法的三元组数量并减去即可。对于一个可能的元素值 \(x\)，值为它的不合法三元组数量为

\[\dbinom{f_x}{3} \]

该组合数可以快速计算。由于值域为 \(\left[1, N\right]\)，故这部分的复杂度也为 \(\mathcal{O}(N)\)。

总复杂度 \(\mathcal{O}(N)\)，可以通过本题。

Code

#include <bits/stdc++.h>

typedef long long valueType;
typedef std::vector<valueType> ValueVector;

valueType C(valueType x) {
    return x * (x - 1) * (x - 2) / 2 / 3;
}

int main() {
    valueType N;

    std::cin >> N;

    ValueVector A(N + 1, 0);
    ValueVector F(N + 1, 0), G(N + 1, 0);

    valueType ans = 0;

    for (valueType i = 1; i <= N; ++i) {
        std::cin >> A[i];

        ans += F[A[i]] * (i - 1) - G[A[i]];

        ++F[A[i]];
        G[A[i]] += i;
    }

    for (valueType i = 1; i <= N; ++i)
        if (F[i] >= 3)
            ans -= C(F[i]);

    std::cout << ans << std::endl;
}