【力扣】480.滑动窗口中位数
题目链接:https://leetcode.cn/problems/sliding-window-median/
面试的时候不会手撕treap,原来两个multiset就能搞定!我是笨蛋。
代码:
class Solution {
public:
multiset<long long> lst, rst;
int lcnt = 0, rcnt = 0;
void insert(long long x) {
if (lcnt == 0 && rcnt == 0)
lst.insert(x), lcnt++;
else {
long long r = *(prev(lst.end()));
if (x <= r)
lst.insert(x), lcnt++;
else
rst.insert(x), rcnt++;
}
if (lcnt - rcnt > 1) {
auto rit = prev(lst.end());
long long tmp = *rit;
lst.erase(rit);
rst.insert(tmp);
lcnt--, rcnt++;
} else if (rcnt > lcnt) {
auto lit = rst.begin();
long long tmp = *lit;
rst.erase(lit);
lst.insert(tmp);
lcnt++, rcnt--;
}
}
void de(long long x) {
long long r = *(prev(lst.end()));
if (x <= r) {
auto it = lst.find(x);
lst.erase(it);
lcnt--;
} else {
auto it = rst.find(x);
rst.erase(it);
rcnt--;
}
if (lcnt - rcnt > 1) {
auto rit = prev(lst.end());
long long tmp = *rit;
lst.erase(rit);
rst.insert(tmp);
lcnt--, rcnt++;
} else if (rcnt > lcnt) {
auto lit = rst.begin();
long long tmp = *lit;
rst.erase(lit);
lst.insert(tmp);
lcnt++, rcnt--;
}
}
double get_mid() {
if (lcnt > rcnt) {
auto rit = prev(lst.end());
return 1.0 * (*rit);
} else if (rcnt > lcnt){
auto lit = rst.begin();
return 1.0 * (*lit);
} else {
auto rit = prev(lst.end()), lit = rst.begin();
return 0.5 * ((*rit) + (*lit));
}
}
vector<double> medianSlidingWindow(vector<int> &nums, int k) {
for (int i = 0; i < k; i++) {
insert(nums[i]);
}
vector<double> res;
res.push_back(get_mid());
printf("%lf ", get_mid());
for (int i = k; i < (int)nums.size(); i++) {
insert(nums[i]);
de(nums[i - k]);
res.push_back(get_mid());
printf("%lf ", get_mid());
}
return res;
}
};