【力扣】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;
    }
};
posted @ 2024-03-31 15:27  tudouuuuu  阅读(1)  评论(0编辑  收藏  举报