leetcode刷题笔记-LCP24-set维护中位数

题目链接:https://leetcode-cn.com/problems/5TxKeK/submissions/

题意

给定一个n长度数组,求对于每个$1\le i\le n$将初始数组$0到i-1位变成x,x+1,x+2,...x+i$的最小操作数,每次操作可以将数组的某一位+1或-1

题解

考虑一开始以第一位为基准然后所有数值减小$nums[0]+i$,现在题目变成将所有数值变成同一个数的最小操作数
考虑这个数为基准线,移动基准线时显然朝着数值多的一边移动,所有数离基准线的距离和最小,所以取中位数为基准线就是最小答案
问题变为维护中位数,用两个$set:st,ed$分别维护比中位数大和比中位数小的数,维护答案即可
也可以使用值域线段树/堆维护解题

code

typedef long long ll;
const ll mod = 1000000007;
class Solution
{
public:
    vector<int> numsGame(vector<int> &nums)
    {
        multiset<int> st, ed;
        multiset<int>::iterator it;
        int n = nums.size();
        for (int i = 1; i < n; i++)
            nums[i] -= nums[0] + i;
        int up = 0, down = 0;
        ll curv = 0;
        int pval = 0;
        vector<int> res(n, 0);
        for (int i = 1; i < n; i++)
        {
            int val = nums[i];
            if (val > pval)
            {
                up++;
                ed.insert(val);
            }
            else
            {
                down++;
                st.insert(val);
            }
            curv += abs(pval - val);
            if (up > down + 1)
            {
                it = ed.begin();
                curv -= (*it - pval);
                up--;
                down++;
                st.insert(pval);
                pval = *it;
                ed.erase(it);
            }
            else if (down > up + 1)
            {
                it = st.end();
                it--;
                curv -= (pval - *it);
                ed.insert(pval);
                pval = *it;
                st.erase(it);
                up++;
                down--;
            }
            res[i] = curv % mod;
        }
        return res;
    }
};