HDU 6058 - Kanade's sum | 2017 Multi-University Training Contest 3

/*
HDU 6058 - Kanade's sum [ 思维,链表 ]  |  2017 Multi-University Training Contest 3
题意：
	给出排列 a[N]，求所有区间的第k大数之和
	N <= 5e5, k <= 80
分析：
	先求出每个数的位置pos[]数组
	然后维护一个单调链表，按数字从大到小向里面加该数字的位置
		即对于i的位置 pos[i] 找到链表中比它大的第一个位置和比它小的第一个位置，将pos[i]加到两个之间
	这部分用set实现，复杂度 O(nlogn)
 
	则对于数字i,链表中所有的位置均为比i大的数的位置
	然后滑动窗口更新答案
		即维护 l,r，要求 l 到 r 之间包含 pos[i] 且恰好 k 个数，这个时候 i 就是第k小
	复杂度 O(nk)
 
	总复杂度 O(nlogn + nk)
*/
#include <bits/stdc++.h>
using namespace std;
const int N = 5e5+5;
int t, n, k;
int a[N], pos[N];
int pre[N], nxt[N];
set<int> st;
set<int>::iterator it;
long long ans;
void solve(int x)
{
    int l, r, lx, rx, cnt = k-1;
    l = r = x;
    while (cnt && pre[l] != 0) l = pre[l], cnt--;
    while (cnt && nxt[r] != n+1) r = nxt[r], cnt--;
    while (l != nxt[x])
    {
        lx = l-pre[l];
        rx = nxt[r]-r;
        ans += (long long)lx * rx * a[x];
        if (nxt[r] == n+1) break;
        l = nxt[l];
        r = nxt[r];
    }
}
int main()
{
    scanf("%d", &t);
    while (t--)
    {
        st.clear();
        scanf("%d%d", &n, &k);
        for (int i = 1; i <= n; i++)
        {
            scanf("%d", a+i);
            pos[a[i]] = i;
        }
        pre[n+1] = 0;
        nxt[0] = n+1;
        st.insert(0);
        st.insert(n+1);
        ans = 0;
        for (int i = n; i >= 1; i--)
        {
            it = st.lower_bound(pos[i]);
            nxt[pos[i]] = *it;
            pre[*it] = pos[i];
            --it;
            pre[pos[i]] = *it;
            nxt[*it] = pos[i];
            if (n-i+1 >= k) solve(pos[i]);
            st.insert(pos[i]);
        }
        printf("%lld\n", ans);
    }
}