k大子段和 二分答案 + 平衡树

问题描述：
    给定一个整数序列a[1..N]，定义sum[i][j]=a[i]+a[i+l]+……+a[j]，将所有的sum[i][j]从小到大排序（其中i，j满足1<=i<=j<=N），得到一个长为N*(N+1)/2的序列，求该序列中的第k个元素。
输入格式(ktm.in)
    第一行有两个整数N，k，其中0<N<=20000，1<=k<=N*(N+1)/2，数据保证任何一个sum[i][j]的绝对值不超过2^30。
    接下来N行每行一个整数。顺序给出序列a的元素。
输出格式(kth.out)
    sum序列中的第k个元素
输入样例
5 15
1
2
3
4
5
输出样例
15
30%的数据满足n<=1000；
另外30%的数据满足n<=20000,k<=100000并且所有元素都为正整数。

这道题我一开始受 ccl 的影响，以为是 k 短路问题，但是看到 k 的范围瞬间觉得不科学。

接着 lyp 跟我讲了二分答案的做法，觉得非常厉害。

也就是说，每二分出一个值，便到所有的区间和中查找他的 rank，直到出解为止。

但是还有一个问题：序列中存在负数，也就是说区间和并不是单调的的。那么如何查找 rank 呢？

我的实现是使用平衡树(treap)，每次查找的复杂度为 O(logn)。

这样一来，二分答案复杂度为 O(log(1 << 31))， 枚举区间端点复杂度为 O(n)，查找为 O(logn)，总复杂度为 O(nlogn2)。

总结：又学会了一种问题的解决方法。

Code :

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <climits>
#include <iostream>
#include <algorithm>

typedef int MAINTYPE;
typedef long long int64;
typedef unsigned int uint;
typedef unsigned long long uint64;

#define swap(a, b, t) ({t _ = (a); (a) = (b); (b) = _;})
#define MAX(a, b, t) ({t _ = (a), __ = (b); _ > __ ? _ : __;})
#define MIN(a, b, t) ({t _ = (a), __ = (b); _ < __ ? _ : __;})

#define max(a, b) MAX(a, b, MAINTYPE)
#define min(a, b) MIN(a, b, MAINTYPE)

#define maxn 200105
#define random (rand() * rand())
#define sum(i, j) (s[j] - s[(i) - 1])

int64 n, k;
int a[maxn], s[maxn];
int mini, maxi, ans, vetot;
struct node{node * c[2]; int data, rank, size;} vess[maxn], * root;

void update(node * & p)
{
    if (! p) return;
    p->size = (p->c[0] ? p->c[0]->size : 0) + (p->c[1] ? p->c[1]->size : 0) + 1;
}

void rotate(node * & p, bool flag)
{
    node * q = p->c[! flag];
    p->c[! flag] = q->c[flag], q->c[flag] = p;
    update(p), update(q), p = q;
}

void insert(node * & p, int data)
{
    if (! p)
    {
        p = vess + ++ vetot;
        p->data = data, p->rank = random, p->size = 1;
        p->c[0] = p->c[1] = 0;
    }
    else
    {
        bool flag = data > p->data;
        insert(p->c[flag], data), ++ p->size;
        if (p->c[flag]->rank < p->rank) rotate(p, ! flag);
    }
}

int getrank(int data)
{
    node * p = root;
    int rank = 0;
    while (p)
    {
        rank += p->c[0] ? p->c[0]->size + 1: 1;
        if (data <= p->data)
        {
            rank -= p->c[0] ? p->c[0]->size + 1: 1;
            p = p->c[0];
        }
        else
            p = p->c[1];
    }
    return rank;
}

void work()
{
    if (k == 1) return (void) printf("%d\n", mini);
    if (k == n * (n + 1) / 2) return (void) printf("%d\n", maxi);
    for (int l = mini, r = maxi; l < r; )
    {
        int mid = (l + r) >> 1, rank = 0;
        vetot = 0, root = 0;
        for (int i = n; i >= 1; -- i)
        {
            insert(root, s[i]);
            rank += getrank(mid + s[i - 1]);
        }
        rank < k ? (l = mid + 1) : (r = mid);
        ans = l - 1;
    }
    printf("%d\n", ans);
}

void prepare()
{
    mini = INT_MAX, maxi = INT_MIN;
    for (int i = 1, s = 0, t = 0; i <= n; ++ i)
    {
        mini = min(mini, s += a[i]);
        maxi = max(maxi, t += a[i]);
        if (s > 0) s = 0;
        if (t < 0) t = 0;
    }
}

void init()
{
    scanf("%I64d%I64d", & n, & k);
    for (int i = 1; i <= n; ++ i)
        scanf("%d", & a[i]), s[i] = s[i - 1] + a[i];
}

int main()
{
    freopen("kth.in", "r", stdin);
    freopen("kth.out", "w", stdout);
    
    init();
    prepare();
    work();
    
    return 0;
}