给定长度为N的序列,对于从前往后每个长度为M的子串求出字串中前K小数的和
解法一
先离散化数据,然后将滑动窗口内的数字个数插入树状数组,数字的值的和插入另一个树状数组,
每次对窗口内的元素查询,二分枚举第一个树状数组前 mid个位置是否存在k个数,查询到包含有k个数的位置后,直接查询另一个树状数组的该位置之前的数字的和即可。要注意,我们不一定能找到一个位置恰好在它之前有k个数,所以要特殊处理下当前有效位置包含的数小于k,下一个有效位置包含的数大于k的情况。详见代码
#include <bits/stdc++.h>
#define int long long
#define endl '\n'
#define sz(x) x.size()
#define lbt(x) (x)&(-x)
#define rep(i,n) for(int i=0;i<n;i++)
using namespace std;
typedef pair<int, int> PII;
typedef double db;
const int N = 2e5 + 10, mod = 998244353, inf = 0x3f3f3f3f;
const db EPS = 1e-9;
int n, m, k;
int t[2][N];
vector<int>v, b;
map<int, int>id;
void add(int u, int i, int x) {
while (i <= n) {
t[u][i] += x;
i += lbt(i);
}
}
int ask(int u, int i) {
int res = 0;
while (i > 0) {
res += t[u][i];
i -= lbt(i);
}
return res;
}
int query() {
int l = 1, r = n;
while (l < r) {
int mid = l + r + 1 >> 1;
if (ask(0, mid) <= k)l = mid;
else r = mid - 1;
}
if (ask(0, l) == k)return ask(1, l);
return ask(1, l) + (k - ask(0, l)) * b[l];
}
void solve() {
cin >> n >> m >> k;
v.resize(n + 1);
v[0] = 2e9 + 10;
for (int i = 1;i <= n;i++)cin >> v[i];
b = v;
sort(b.begin(), b.end());
b.erase(unique(b.begin(), b.end()), b.end());
for (int i = 0;i < b.size();i++)id[b[i]] = i + 1;
for (int i = 1;i <= n;i++) {
add(0, id[v[i]], 1);
add(1, id[v[i]], v[i]);
if (i >= m) {
if (i > m) {
add(0, id[v[i - m]], -1);
add(1, id[v[i - m]], -v[i - m]);
}
cout << query() << " ";
}
}
}
signed main() {
cin.tie(0)->sync_with_stdio(0);
cout.tie(0);
int _ = 1;
// cin >> _ ;
while (_--)
solve();
return 0;
}
解法二
用两个set维护下整个窗口,前k小的在第一个set(in),后面的部分在另一个se(out)t中,维护过程保证in.size()<k,细节很多,模拟时逻辑一定要清晰。
#include <bits/stdc++.h>
#define int long long
#define endl '\n'
#define sz(x) x.size()
#define lbt(x) (x)&(-x)
#define rep(i,n) for(int i=0;i<n;i++)
using namespace std;
typedef pair<int, int> PII;
typedef double db;
const int N = 2e5 + 10, mod = 998244353, inf = 0x3f3f3f3f;
const db EPS = 1e-9;
int a[N];
void solve() {
int n, m, k;
cin >> n >> m >> k;
int res = 0;
multiset<int>in, out;
for (int i = 1;i <= n;i++) {
cin >> a[i];
if (i <= m) {
res += a[i];
in.insert(a[i]);
if (in.size() > k) {
res -= *in.rbegin();
out.insert(*in.rbegin());
in.erase(in.find(*in.rbegin()));
}
}
else {
int w = a[i - m];
in.insert(a[i]);
res += a[i];
res -= *in.rbegin();
out.insert(*in.rbegin());
in.erase(in.find(*in.rbegin()));
if (w <= *in.rbegin()) {
res -= w;
res += *out.begin();
in.insert(*out.begin());
out.erase(out.find(*out.begin()));
in.erase(in.find(w));
}
else out.erase(out.find(w));
}
if (i >= m)cout << res << " ";
}
}
signed main() {
cin.tie(0)->sync_with_stdio(0);
cout.tie(0);
int _ = 1;
// cin >> _ ;
while (_--)
solve();
return 0;
}