S - 数据结构复习 E. 第K大和

题目链接

妙妙题！

难度：Medium-。

题意

给定一个 $1\sim n$ 的全排列 $A_1,A_2,\cdots,A_n$。给定一个 $k$，统计所有长度 $\geq k$ 的子区间的第 $k$ 大的数的和。

$n\leq 5\times 10^5$，$k\leq \min(n,80)$。

题解

考虑如果 $k=1$ 怎么做。

显然算每个数对答案的贡献，然后求出左右两边第一个比 $A_i$ 大的位置 $L$ 和 $R$，则 $A_i$ 对答案的贡献为 $(i-L)\times (R-i)$。

考虑拓展到 $k>1$ 的做法。

可以从小到大考虑 $A_i$ 的值，每次算完贡献就把 $A_i$ 删掉，这样数列中剩下的就都是比 $A_i$ 大的数了。

如果可以求出 $i$ 往前 $k$ 个比 $A_i$ 大的位置和后面 $k$ 个比 $A_i$ 大的位置的话，就可以 $\mathcal{O}(k)$ 算出答案了。

这个可以用链表维护。时间复杂度 $\mathcal{O}(nk)$。

代码

#include <bits/stdc++.h>
#define int long long
#define rep(i, j, k) for (int i = (j); i <= (k); i++)
#define per(i, j, k) for (int i = (j); i >= (k); i--)
#define ll long long
#define vi vector<int>
#define sz(a) ((int)(a.size()))
#define mset(f, x) memset(f, x, sizeof(f))
#define ALL(x) (x).begin(), (x).end()
#define rALL(x) (x).rbegin(), (x).rend()
#define uni(x) x.resize(unique(ALL(x)) - x.begin())
#define ull unsigned long long
#define pii pair<int, int>
using namespace std;
inline int min(int x, int y) { return x < y ? x : y; }
inline int max(int x, int y) { return x > y ? x : y; }
inline void chmin(int &x, int y)
{
    if (y < x)
        x = y;
}
inline void chmax(int &x, int y)
{
    if (y > x)
        x = y;
}
const int N = 5e5 + 10;
int n, k, a[N], pos[N];
int pre[N], nxt[N];
vector<int> vecL, vecR;
void Main()
{
    cin >> n >> k;
    rep(i, 1, n)
    {
        cin >> a[i];
        pos[a[i]] = i;
        pre[i] = i - 1;
        nxt[i] = i + 1;
    }
    int ans = 0;
    rep(x, 1, n)
    {
        int p = pos[x];
        vecL.clear();
        vecR.clear();
        int u = p;
        rep(t, 1, k + 1)
        {
            vecL.push_back(u);
            u = pre[u];
            if (u < 1)
            {
                if (t <= k)
                    vecL.push_back(u);
                break;
            }
        }
        u = p;
        rep(t, 1, k + 1)
        {
            vecR.push_back(u);
            u = nxt[u];
            if (u > n)
            {
                if (t <= k)
                    vecR.push_back(u);
                break;
            }
        }
        rep(l, 1, k)
        {
            int r = k + 1 - l;
            if (l < sz(vecL) && r < sz(vecR))
                ans += x * (vecL[l - 1] - vecL[l]) * (vecR[r] - vecR[r - 1]);
        }
        // cerr << ans << endl;
        pre[nxt[p]] = pre[p];
        nxt[pre[p]] = nxt[p];
    }
    cout << ans << "\n";
}
signed main()
{
    ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
    int t = 1;
    cin >> t;
    while (t--)
        Main();
    return 0;
}