2020牛客寒假算法基礎集訓營4 H題

坐火車

題目描述

牛牛是一名喜歡旅遊的同學，在來到渡渡鳥王國時，坐上了顏色多樣的火車。
牛牛同學在車上，車上有 n 個車廂，每一個車廂有一種顏色。
他想知道對於每一個正整數x∈[1, n] ，集合{(i, x, j) ∣ i<x<j, l x≤col i=col j≤rx}中包含多少個元素。
換句話說，就是要求每一個車廂兩邊有多少對顏色相同的車廂，並且這一對車廂的顏色要在 lx 到 rx 之間。其中 coli 代表 i 號車廂的顏色，lx,rx 代表顏色的限制。

輸入描述:

第一行一個正整數n。
第二行 n 個三元組，每個三元組包括三個正整數 (coli, li, ri)，輸入中沒有括號，這 3n 個正整數之間均只用空格隔開，詳見樣例。

輸出描述:

輸出一行 n 個非負整數代表答案。

---

這道題思維比較難想；

解釋一下題解代碼：

思考的對象可以放在每次枚舉的車廂的顏色上，從1–n枚舉每個車廂，當前車廂的顏色爲x，這個位置的車廂就不能與前面的車廂顏色爲x的匹配，所有我們要減去當前位置 i 前面x的個數（因爲我們之前就已經加上），還要加上當前位置 i 後面的 x 的個數（意味着這個位置 i 可以和他們匹配）；

所以樹狀數組就是一個不錯的選擇；

#include<bits/stdc++.h>
#define ll long long
#define pa pair<int,int>
#define lson k<<1
#define rson k<<1|1
//ios::sync_with_stdio(false);
using namespace std;
const int N=500100;
const int M=200100;
const ll mod=998244353;
ll s1[N],s2[N],tr[N];//s1代表數i的後面還有幾個數i，s2代表數i前面還有幾個數i 
int n;
struct Node{
	int l,r,c;
}col[N];
int lowbit(int k){
	return k & (-k);
}
void add(int p,ll q){//p點加上q，單點修改 
	while(p<=N){
		tr[p]+=q;
		p+=lowbit(p);
	}
}
ll sum(int p){//單點查詢，就是求前綴和 
	ll s=0;
	while(p!=0){
		s+=tr[p];
		p-=lowbit(p);
	}
	return s;
}
int main(){
    ios::sync_with_stdio(false);
    cin>>n;
    for(int i=1;i<=n;i++) cin>>col[i].c>>col[i].l>>col[i].r;
    for(int i=n;i>=1;i--) s1[col[i].c]++;
    for(int i=1;i<=n;i++){
    	s1[col[i].c]--;//後面還有幾個col[i].c 
    	add(col[i].c,-s2[col[i].c]);//減去前面col[i].c的個數 
    	cout<<sum(col[i].r)-sum(col[i].l-1)<<" ";
    	s2[col[i].c]++;//前面還有幾個col[i].c
    	add(col[i].c,s1[col[i].c]); 
	} 
    return 0;
}