2024美团春季笔试第二场 T4 区间众数求和 线段树做法
(以下讲解的下标均从1开始)
题目大意
给出一个数组 a[1~n]，a[i] = 1或2， 求每个区间的众数之和。
思路
写公式
我们假设 two[i] 为 a[1~i] 的“2”的数量，one[i] 为 a[1~i] 的“1”的数量，假设m为众数为2的区间数量，那么有下面的公式：
其中[...]表示如果括号内的条件满足，则=1，否则=0.
这个式子可以转换为：
设一个数组 diff[i] = two[i] - one[i]。那么式子再写成：
可以看到这个式子非常的熟悉，可能有人会想到一维偏序什么的。
首先，-n <= diff[i] <= n。我们不喜欢负数的存在，由于比较大小时两边都加同一个数没啥毛病，考虑每个diff[i]都 +(n+1)，这样就都是正数了。
回到正题，如果暴力求解这个式子的话，显然要用O(n^2)的查询。我们希望遍历r时，可以马上获取数组中在r左边，且 < diff[r]的数的数量，这里，我当时就用的线段树。
请出线段树
（不知道线段树的朋友可以先去了解一下）
建立一个线段树tree，tree[i] 表示树的第i号节点，我们先假设其表示的区间为[left,right]。则tree[i]的值就是diff[1~r-1]的所有值处于[left,right]的元素数量。比如diff[] = {3,5,4,7,1, ... }，[left,right] = [1,4]，假设r此时为5.则tree[i]就是3,5,4,7（下标<r）中位于[1,4]的数：3,4，它们数量的和：2。则tree[i] = 2.
假设我们处理了diff[1~r-1]，把它们写到了线段树里，现在轮到diff[r]了。
我们要判断diff[1~r-1]中有多少<diff[r]的，怎么办？我们就去线段树中去找在区间[left,right] = [0,diff[r]-1]内的tree值。这个[0,diff[r]-1]是可以用线段树的性质拼接出来的，对每个组成其一部分的子区间的tree值求和即可得到答案，对应的增加众数为2的区间数量的贡献。
查找完之后我们还要将diff[r]的值加到线段树里，这样在继续找diff[r+1,n]的答案时可以把diff[r]的贡献算进去。
最后，由于总区间个数为sum = n*(n+1)/2，众数为2的区间个数为m，则众数为1的区间个数为p = sum - m，答案为m*2 + p。
复杂度
时间复杂度O(nlogn)，其中线段树查询和单点修改为O(logn)，做n次
空间复杂度O(n)
#include <iostream>
#include <set>
#include <algorithm>
#include <vector>
#include <cstring>
using namespace std;
using ll = long long;
const int N = (int)2e5+5;
ll n;
ll a[200005];
ll tree[N << 3]; // 开大点没关系
ll diff[200005];
// 将数值x放入线段树，此时所在的区间为[l,r]，线段树下标为id
void push(ll x,int l,int r,int id){
if(l >= r){
tree[id]++;
return;
}
int mid = (l+r)>>1;
if(x <= mid){
push(x,l,mid,id<<1);
}
else{
push(x,mid+1,r,(id<<1)+1);
}
tree[id] = tree[id<<1] + tree[(id<<1)+1];
}
// 求<x的diff数量
ll getsum(ll x,int l,int r,int id){
if(l >= r){
if(x <= l) return 0;
return tree[id];
}
ll res = 0;
int mid = (l+r)>>1;
if(x > mid){
res += tree[id<<1];
res += getsum(x,mid+1,r,(id<<1)+1);
}
else{
res += getsum(x,l,mid,(id<<1));
}
return res;
}
void solve(){
ll m = 0;
for(int i = 1;i <= n;++i){
cin >> a[i]; // 1 <= a[i] <= 2
diff[i] = diff[i-1] + (a[i] == 2 ? 1 : -1);
}
for(int i = 0;i <= n;++i){
diff[i] += n+1; // 保证所有数为正数
}
push(diff[0],0,2*n+1,1); // 由于l-1可以为0，因此先把diff[0]放进去
for(int i = 1;i <= n;++i){
m += getsum(diff[i],0,2*n+1,1); // 从最大的区间[0,2n+1]开始递归查询
push(diff[i],0,2*n+1,1); // 单点修改，其所在的区间的值都要+1
}
ll p = n*(n+1)/2 - m;
ll ans = m*2 + p;
cout << ans << "\n";
}
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
cout.tie(nullptr);
while(cin >> n){
solve();
}
return 0;
}
投递美团等公司10个岗位 >