【Luogu P7527】[USACO21OPEN] United Cows of Farmer John G（树状数组）

题目

题目描述

农夫约翰合牛国（The United Cows of Farmer John，UCFJ）将要选派一个代表队参加国际牛学奥林匹克（International bOvine olympIad，IOI）。

有 $N$ 头奶牛参加了代表队选拔。她们站成一行，奶牛 $i$ 的品种为 $b_i$。

代表队将会由包含至少两头奶牛的连续区间组成——也就是说，对于满足 $1\le l<r\le N$ 的奶牛 $l\dots r$。最边上的奶牛会被指定为领队。为了避免种内冲突，每一名领队都必须与代表队的其他成员（包括领队）品种不同。

请帮助 UCFJ 求出他们可以选派参加 IOI 的代表队的方法数。

输入格式

输入的第一行包含 $N$。

第二行包含 $N$ 个整数 $b_1,b_2,\dots ,b_N$，均在范围 $[1,N]$ 之间。

输出格式

输出一行表示可能的代表队的数量。

输入输出样例

输入 #1
7
1 2 3 4 3 2 5
输出 #1
13

说明/提示

样例解释
每一代表队对应以下的一对领队：
$(1,2),(1,3),(1,4),(1,7),(2,3),(2,4),(3,4),(4,5),(4,6),(4,7),(5,6),(5,7),(6,7)$.

数据范围与约定
$1\le N\le 2\times 10^5$。

思路

记 $l_i$ 为第 $i$ 头牛作为右端点，最多能往左到哪里。
记 $r_i$ 为第 $i$ 头牛作为左端点，最多能往右到哪里。
$l$ 和 $r$ 数组可以轻易求出，正、反各扫一遍即可。
考虑 $(a,b)$ 何时满足条件，容易发现只要当 $l_b\le a$ 并且 $r_a\ge b$ 即可。
那么，从左往右扫，对于第 $i$ 头牛，只需要求出有多少个 $l_j\le i$，其中 $j\in [i+1,r_i]$，将个数累加即可。
于是问题就被转化成了“区间内 $\le k$ 的数的个数”。
有很多方法求解，这里采用树状数组。
记得开 long long。

代码

#include <bits/stdc++.h>
#define rep(i, a, b) for (int i = a; i <= b; i++)
#define per(i, a, b) for (int i = a; i >= b; i--)
#define mem(a, x) memset(a, x, sizeof a)
#define pb push_back
#define umap unordered_map
#define pqueue priority_queue
#define PI acos(-1)
#define int long long

using namespace std;
typedef long long ll;
int l[200010], r[200010], a[200010], h[200010];
int n, tr[200010];
struct node1 {
	int x, id;
} A[200010];
struct node2 {
	int L, R, x;
} b[200010];

template <typename _T>
void rd(_T &x) {
    int f = 1; x = 0;
    char s = getchar();
    while (s > '9' || s < '0') {if (s == '-') f = -1; s = getchar();}
    while (s >= '0' && s <= '9') x = (x<<3)+(x<<1)+(s-'0'), s = getchar();
    x *= f;
}

inline bool cmp1(const node1 &x, const node1 &y) {
	return x.x < y.x;
}

inline bool cmp2(const node2 &x, const node2 &y) {
	return x.x < y.x;
}

inline int lowbit(int x) {
	return x & (-x);
}

void add(int x, int b) {
	while (x <= n) {
		tr[x] += b;
		x += lowbit(x);
	}
	return ;
}

int query(int x) {
	int ret = 0;
	while (x != 0) {
		ret += tr[x];
		x -= lowbit(x);
	}
	return ret;
}

signed main() {
	rd(n);
	rep(i, 1, n) rd(a[i]);
	rep(i, 1, n) l[i] = h[a[i]]+1, h[a[i]] = i;
	rep(i, 1, n) h[i] = n+1;
	per(i, n, 1) r[i] = h[a[i]]-1, h[a[i]] = i;
	rep(i, 1, n) A[i] = (node1){l[i], i};
	sort(A+1, A+1+n, cmp1);
	rep(i, 1, n) b[i] = (node2){i+1, r[i], i};
	sort(b+1, b+1+n, cmp2);
	int ans = 0, j = 1;
	rep(i, 1, n) {
		for (; j <= n && A[j].x <= b[i].x; j++) add(A[j].id, 1);
		ans += query(b[i].R)-query(b[i].L-1);
	}
	printf("%lld", ans);
	return 0;
}