SGU 311. Ice-cream Tycoon（平衡树）

题目链接：点击打开链接

思路：

最简单的思路是， 维护一棵平衡树， 以单价作为键值， 维护一个结点个数的信息。  然后类似寻找第K小的方法找到第n个小的价值处， 顺便维护总价值， 判断是否happy。

然后就是删除操作。 我们从新从根结点开始， 沿路删除可以删除的子树就行了。 由于每次删除都是删除的叶子结点处， 所以不需要旋转操作。

细节参见代码：

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <iostream>
#include <string>
#include <vector>
#include <stack>
#include <bitset>
#include <cstdlib>
#include <cmath>
#include <set>
#include <list>
#include <deque>
#include <map>
#include <queue>
#define Max(a,b) ((a)>(b)?(a):(b))
#define Min(a,b) ((a)<(b)?(a):(b))
using namespace std;
typedef long long ll;
typedef long double ld;
const ld eps = 1e-9, PI = 3.1415926535897932384626433832795;
const int mod = 1000000000 + 7;
const int INF = 0x3f3f3f3f;
const int seed = 131;
const ll INF64 = ll(1e18);
ll n, c;
struct node {
    node *ch[2];
    int r;      // 左右子树
    ll s, tt, t, ss;   //ss是该结点的物品数，s是结点总数，tt是单价，t是总价， r是优先级
    node(ll tt=0, ll ss=0): tt(tt), ss(ss) {
        ch[0] = ch[1] = NULL;
        s = ss;
        t = tt * ss;
        r = rand();
    }
    ll cmp(ll x) const {
        if(x == tt) return -1;
        return x < tt ? 0 : 1;
    }
    void maintain() {
        s = ss;
        t = tt * ss;
        if(ch[0] != NULL) s += ch[0]->s, t += ch[0]->t; //维护这个结点下的所有节点数， 维护一个结点下的总价钱
        if(ch[1] != NULL) s += ch[1]->s, t += ch[1]->t;
    }
} *g;
void rotate(node* &o, ll d) {
    node* k = o->ch[d^1];  //旋转， 使得优先级满足堆的意义
    o->ch[d^1] = k->ch[d];
    k->ch[d] = o;
    o->maintain();
    k->maintain();
    o = k;
}
void insert(node* &o, ll x, ll s) {
    if(o == NULL) o = new node(x, s);
    else {
        int d = (x < o->tt ? 0 : 1);
        insert(o->ch[d], x, s);
        if(o->ch[d]->r > o->r) rotate(o, d^1);
    }
    o->maintain();
}
void removetree(node* &x) {
    if(x->ch[0] != NULL) removetree(x->ch[0]);
    if(x->ch[1] != NULL) removetree(x->ch[1]);
    delete x;
    x = NULL;
}
void _remove(node* &gg, node* &o, ll k, ll sum) {
    if(o == NULL || k <= 0 || k > o->s) return ;
    ll s = (o->ch[0] != NULL ? o->ch[0]->s : 0); //s是左子树的元素个数
    ll hehe = (o->ch[0] != NULL ? o->ch[0]->t : 0);
    if(k <= s) {
        _remove(o, o->ch[0], k, sum);
        o->maintain();
    }
    else if(k <= s + (o->ss)) {
        ll tot = sum; //这k件最便宜的物品的总价
        if(o->ch[0] != NULL) tot += o->ch[0]->t;
        tot += o->tt * (k - s);
            o->s -= k;
            o->ss -= (k - s);
            o->t -= tot - sum;
            if(o->ch[0] != NULL) removetree(o->ch[0]);
    }
    else {
        ll hehe2 = (o->ss * o->tt);
        ll hehe3 = o->ss;
        node* u = o;
        if(o->ch[0] != NULL) removetree(o->ch[0]);
        o = o->ch[1];
        delete u;
        _remove(gg, o, k-s-hehe3, sum + hehe + hehe2);
        o->maintain();
    }
    o->maintain();
}
bool solve(node* o, ll k, ll sum) {
    if(o == NULL || k <= 0 || k > o->s) return false;
    ll s = (o->ch[0] != NULL ? o->ch[0]->s : 0); //s是左子树的元素个数
    ll hehe = (o->ch[0] != NULL ? o->ch[0]->t : 0);
    if(k <= s) {
        int ok = solve(o->ch[0], k, sum);
        o->maintain();
        return ok;
    }
    else if(k <= s + (o->ss)) {
        ll tot = sum; //这k件最便宜的物品的总价
        if(o->ch[0] != NULL) tot += o->ch[0]->t;
        tot += o->tt * (k - s);
        if(tot > c) return false;
        else return true;
    }
    else {
        int ok = solve(o->ch[1], k-s-(o->ss), sum + hehe + (o->ss * o->tt));
        o->maintain();
        return ok;
    }
    o->maintain();
}
char s[100];
int main() {
    g = NULL;
    while(~scanf("%s%I64d%I64d",s, &n, &c)) {
        if(s[0] == 'A') {
            insert(g, c, n);
        }
        else {
            if(solve(g, n, 0)) _remove(g, g, n, 0), printf("HAPPY\n");
            else printf("UNHAPPY\n");
        }
    }
    return 0;
}