AT_abc318_e Sandwiches 题解
Links
Description
给定一个长度为 \(n\) 的序列 \(a\)，找到满足以下条件的三元组 \((a,b,c)\) 的数量。
- \(i < j < k\)；
- \(a_{i} = a_{k}\)；
- \(a_{i} \neq a_{j}\)。
数据范围：\(1 \leq n \leq 3 \times 10^{5}\)，\(1 \leq a_{i} \leq n\)。
Solution
由于要求 \(i < j < k\)，我们可以从左向右枚举 \(k\)，考虑对于这个 \(k\) 计算贡献。假设 \(a_{k}\) 在之前出现次数是 \(cnt_{a_{k}}\)，出现的位置集合是 \(pos_{a_{k}} = \{pos_{a_{k},1}, \cdots, pos_{a_{k},cnt_{a_{k}}}\}\)。先考虑一种较劣的算法：枚举每个 \(i\)，那么 \(j\) 的数量可以很轻松的确定（注意减去中间 \(a_{j}\) 和 \(a_{k}\) 相同的数量）。于是这个 \(k\) 的贡献是：
\[\sum_{i = 1}^{cnt_{a_{k}}} \left ( k - pos_{a_{k},i} - cnt_{a_{k}} + i \right) \]此时复杂度为 \(\Omicron (n^{2})\)。
考虑拆上面的式子，令其为 \(f(k)\)：
\[\begin{aligned} f(k) &= \sum_{i = 1}^{cnt_{a_{k}}} k - \sum_{i = 1}^{cnt_{a_{k}}} pos_{a_{k},i} - \sum_{i = 1}^{cnt_{a_{k}}} cnt_{a_{k}} + \sum_{i = 1}^{cnt_{a_{k}}}i \\ &=cnt_{a_{k}} \cdot k - \sum_{i = 1}^{cnt_{a_{k}}} pos_{a_{k},i} - cnt_{a_{k}}^{2} + \frac{cnt_{a_{k}} \times (cnt_{a_{k}} + 1)}{2} \\ &=cnt_{a_{k}} \cdot k - \sum_{i = 1}^{cnt_{a_{k}}} pos_{a_{k},i} - \frac{cnt_{a_{k}} \times (cnt_{a_{k}} - 1)}{2} \end{aligned} \]对每个数字为何一个出现位置的前缀和 \(sum = \sum_{i = 1}^{cnt_{a_{k}}} pos_{a_{k},i}\) 就可以 \(\Omicron(1)\) 计算每个 \(k\) 的贡献了，总时间复杂度 \(\Omicron (n)\)。
Codes
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define max_n 510101
void read(int &p)
{
p = 0;
int k = 1;
char c = getchar();
while (c < '0' || c > '9')
{
if (c == '-')
{
k = -1;
}
c = getchar();
}
while (c >= '0' && c <= '9')
{
p = p * 10 + c - '0';
c = getchar();
}
p *= k;
return;
}
void write_(int x)
{
if (x < 0)
{
putchar('-');
x = -x;
}
if (x > 9)
{
write_(x / 10);
}
putchar(x % 10 + '0');
}
void writesp(int x)
{
write_(x);
putchar(' ');
}
void writeln(int x)
{
write_(x);
puts("");
}
int n,ans = 0,nums[max_n],sum[max_n],cnt[max_n];
signed main()
{
read(n);
for(int i = 1;i <= n;i++)
{
read(nums[i]);
}
for(int i = 1;i <= n;i++)
{
if(cnt[nums[i]] >= 1)
{
ans += cnt[nums[i]] * i - sum[nums[i]] - ((cnt[nums[i]] + 1) * (cnt[nums[i]]) / 2);
}
sum[nums[i]] += i;
cnt[nums[i]]++;
}
writeln(ans);
return 0;
}