Codeforces Round #510 (Div. 2)--D. Petya and Array--（红黑树或树状数组）

 链接：https://codeforces.com/problemset/problem/1042/D                       

                                        D. Petya and Array

Petya has an array aa consisting of nn integers. He has learned partial sums recently, and now he can calculate the sum of elements on any segment of the array really fast. The segment is a non-empty sequence of elements standing one next to another in the array.

Now he wonders what is the number of segments in his array with the sum less than tt . Help Petya to calculate this number.

More formally, you are required to calculate the number of pairs l,r(l≤r ) such that al+al+1+⋯+ar−1+ar<t

Input

The first line contains two integers nn and tt (1≤n≤200000,|t|≤2⋅10^14 ).

The second line contains a sequence of integers a1,a2,…,ana1,a2,…,an (|ai|≤10^9 ) — the description of Petya's array. Note that there might be negative, zero and positive elements.

Output

Print the number of segments in Petya's array with the sum of elements less than tt .

Examples

Input

5 4
5 -1 3 4 -1

Output

5

Input

3 0
-1 2 -3

Output

4

Input

4 -1
-2 1 -2 3

Output

3

Note

In the first example the following segments have sum less than 4 :

• [2,2] , sum of elements is −1
• [2,3] , sum of elements is 2
• [3,3][ , sum of elements is 3
• [4,5], sum of elements is 3
• [5,5] , sum of elements is −1
•  

题意：

给定一个数组，问有多少个不同的区间[l,r] (l<=r)使得区间和小于给定的数t.

解法一（树状数组）：

类似用树状数组求序列逆序对的做法

维护前缀和sum[ i ],所以要求转化为求满足sum[ i ]-sum[ j ] < t的所有（i , j）逆序对是求所有满足a[ i ] < a[ j ] (i > j)的（i,j）对数，

所以将上式转化为sum[ j ] > sum[ i ] - t,具体做法是：可以开两个数组a,b两个数组均记录前缀和，给b排序，b的作用主要是获得a[ i ]-t在b中的排名，从而利用树状数组查询当前有多少前缀和比sum[ i ] - t小，（没看懂的话看代码吧。

#include<bits/stdc++.h>

using namespace std;
typedef long long ll;
const int maxn=200005;
int n,sum[maxn];
ll a[maxn],b[maxn],t;

struct bit{
	int lowbit(int x){
		return x&(-x);
	}
	int getsum(int k){	
		int res=0;
		while(k>=1){
			res+=sum[k];
			k-=lowbit(k);
		}
		return res;
	}
	
	void add(int x,int val){
		while(x<=n+1){
			sum[x]+=val;
			x+=lowbit(x);
		}
	}
	void clear(){
		memset(sum,0,sizeof(sum));
	}
};

int main()
{
	bit tree;
	scanf("%d %lld",&n,&t);
	for(int i=2;i<=n+1;i++) {
		scanf("%lld",&a[i]);
		a[i]+=a[i-1];
		b[i]=a[i];
	}
	sort(b+1,b+n+2);ll ans=0;
	for(int i=1;i<=n+1;i++){
		int rank=upper_bound(b+1,b+n+2,a[i]-t)-b-1;
		ans+=i-tree.getsum(rank)-1;
		int pos=lower_bound(b+1,b+n+2,a[i])-b;
		tree.add(pos,1);
	}
	printf("%lld\n",ans);
 } 

 

解法二（红黑树）：

维护前缀和sum[i],则以i为r的区间个数为所有满足sum[i]-sum[j]<t的j的个数（j<i）.该式变形为-sum[j]<t-sum[i]，所以可以用一颗红黑树维护-sum[j]（点击查看红黑树的库实现博客）,然后用order_of_key查询t-sum[i]的rank,这个rank其实就是右端点为i时左端点满足的个数。总复杂度nlog(n).其中遍历数组为o(n),红黑树查询为log(n).

注：由于红黑树里不能有相同元素，所以这里让红黑树维护pair<-sum[i],index>,其中index为下标i.这样就能保证能存储相同的-sum[i].看红黑树实现起来多简短！

代码：

#include<bits/stdc++.h>
#include<ext/pb_ds/assoc_container.hpp>
 
using namespace std;
using namespace __gnu_pbds;
 
#define pii pair<int, int>
typedef long long ll;
typedef tree<pair<ll,int>,null_type,less< pair<ll,int> >,rb_tree_tag,tree_order_statistics_node_update> rbtree;
 
int n;
ll tmp,t,sum=0,ans=0;
 
int main()
{
	rbtree it;
	scanf("%d%lld",&n,&t);
	it.insert(make_pair(0,0));
	for(int i=1;i<=n;i++){
		scanf("%lld",&tmp);
		sum+=tmp;
		ans+=it.order_of_key(pair<ll,int>(t-sum,0));
		it.insert(pair<ll,int>(-sum,i));
	}
	printf("%lld\n",ans);
}

附上代码实现的红黑树解法：

#include <cstdio>
#include <cctype>
#include <cassert>
using namespace std;

//#define __REDBLACK_DEBUG

#define bro(x) (((x)->ftr->lc == (x)) ? ((x)->ftr->rc) : ((x)->ftr->lc))
#define islc(x) ((x) != NULL && (x)->ftr->lc == (x))
#define isrc(x) ((x) != NULL && (x)->ftr->rc == (x))

template<typename T>
class redblacktree {
    protected:
        struct Node;

        Node* _root;    根节点位置
        Node* _hot; 临时维护的节点

        void init(T);
        void connect34(Node*, Node*, Node*, Node*, Node*, Node*, Node*);
        void SolveDoubleRed(Node*); 双红修正
        void SolveDoubleBlack(Node*);   //双黑修正
        Node* find(T, const int);   允许重复的查找
        Node* rfind(T, const int);  不允许重复的查找
        Node* findkth(int, Node*);
        int find_rank(T, Node*);
#ifdef __REDBLACK_DEBUG
        void checkconnect(Node*);
        void previs(Node*, int);
        void invis(Node*, int);
        void postvis(Node*, int);
#endif

    public:

        struct iterator;

        redblacktree() : _root(NULL), _hot(NULL) {}

        int get_rank(T);
        iterator insert(T);
        bool remove(T);
        int size();
        bool empty();
        iterator kth(int);   //返回第k大的迭代器
        iterator lower_bound(T); 
        iterator upper_bound(T);
#ifdef __REDBLACK_DEBUG
        void vis();
        void correctlyconnected();
#endif
};

template <typename T>
struct redblacktree<T>::Node {
    T val;  节点信息
    bool RBc;   节点颜色，若为true，则节点为Red;否则节点为Black.
    Node* ftr;  父亲
    Node* lc;   左儿子
    Node* rc;   右儿子
    int s;      域

    Node(   T v = T(), bool RB = true,
            Node* f = NULL, Node* l = NULL, Node* r = NULL ,int ss = 1  )
        : val(v), RBc(RB), ftr(f), lc(l), rc(r), s(ss) {}

    Node* succ() {      删除节点时用到的替代节点
        Node* ptn = rc;
        while(ptn->lc != NULL) {
            --(ptn->s);
            ptn = ptn->lc;
        }
        return ptn;
    }

    Node* left_node() {     直接前驱
        Node* ptn = this;
        if(!lc) {
            while(ptn->ftr && ptn->ftr->lc == ptn)
                ptn = ptn->ftr;
            ptn = ptn->ftr;
        } else
            while(ptn->lc)
                ptn = ptn->lc;
        return ptn;
    }

    Node* right_node() {    直接后继
        Node* ptn = this;
        if(!rc) {
            while(ptn->ftr && ptn->ftr->rc == ptn)
                ptn = ptn->ftr;
            ptn = ptn->ftr;
        } else
            while(ptn->rc)
                ptn = ptn->rc;
        return ptn;
    }

    void maintain() {   维护域s
        s = 1;
        if(lc) s += lc->s;
        if(rc) s += rc->s;
    }
};

template <typename T>
struct redblacktree<T>::iterator {
    private:

        Node* _real__node;

    public:

        iterator& operator++() {
            _real__node = _real__node->right_node();
            return *this;
        }

        iterator& operator--() {
            _real__node = _real__node->left_node();
            return *this;
        }

        T operator*() {
            return _real__node->val;
        }

        iterator(Node* node_nn = NULL) : _real__node(node_nn) {}
      //  iterator(T const& val_vv) : _real__node(rfind(val_vv, 0)) {}
        iterator(iterator const& iter) : _real__node(iter._real__node) {}

};

template <typename T>
typename
redblacktree<T>::iterator redblacktree<T>::insert(T v) {
    Node* ptn = find(v, 1);
    if(_hot == NULL) {
        init(v);
        return iterator(_root);
    }
    ptn = new Node(v, true, _hot, NULL, NULL, 1);
    if( _hot->val <= v  )
        _hot->rc = ptn;
    else
        _hot->lc = ptn;
    SolveDoubleRed(ptn);
    return iterator(ptn);
}

template <typename T>
void redblacktree<T>::init(T v) {
    _root = new Node(v, false, NULL, NULL, NULL, 1);
#ifdef __REDBLACK_DEBUG
    ++blackheight;
#endif
}

template <typename T>
typename
redblacktree<T>::Node* redblacktree<T>::find(T v, const int op) {
    Node* ptn = _root;  从根节点开始查找
    _hot = NULL;    维护父亲节点
    while(ptn != NULL) {
        _hot = ptn;
        ptn->s += op;
        if(ptn->val > v)
            ptn = ptn->lc;
        else
            ptn = ptn->rc;
    }
    return ptn;
}

template <typename T>
typename
redblacktree<T>::Node* redblacktree<T>::rfind(T v, const int op) {
    Node* ptn = _root;
    _hot = NULL;
    while(ptn != NULL && ptn->val != v) {
        _hot = ptn;
        ptn->s += op;
        if(ptn->val > v)
            ptn = ptn->lc;
        else
            ptn = ptn->rc;
    }
    return ptn;
}

template <typename T>
void redblacktree<T>::SolveDoubleRed(Node* nn) {
    while((!(nn->ftr)) || nn->ftr->RBc) {
        if(nn == _root) {
            _root->RBc = false;
#ifdef __REDBLACK_DEBUG
            ++blackheight;
#endif
            return;
        }
        Node* pftr = nn->ftr;
        if(!(pftr->RBc)) return;            No double-red
        Node* uncle = bro(nn->ftr);
        Node* grdftr = nn->ftr->ftr;
        if(uncle != NULL && uncle->RBc) {   RR-2
            grdftr->RBc = true;
            uncle->RBc = false;
            pftr->RBc = false;
            nn = grdftr;
        } else {                            RR-1
            if(islc(pftr)) {
                if(islc(nn)) {
                    pftr->ftr = grdftr->ftr;
                    if(grdftr == _root) _root = pftr;
                    else if(grdftr->ftr->lc == grdftr) grdftr->ftr->lc = pftr;
                    else grdftr->ftr->rc = pftr;
                    connect34(pftr, nn, grdftr, nn->lc, nn->rc, pftr->rc, uncle);
                    pftr->RBc = false;
                    grdftr->RBc = true;
                } else {
                    nn->ftr = grdftr->ftr;
                    if(grdftr == _root) _root = nn;
                    else if(grdftr->ftr->lc == grdftr) grdftr->ftr->lc = nn;
                    else grdftr->ftr->rc = nn;
                    connect34(nn, pftr, grdftr, pftr->lc, nn->lc, nn->rc, uncle);
                    nn->RBc = false;
                    grdftr->RBc = true;
                }
            } else {
                if(islc(nn)) {
                    nn->ftr = grdftr->ftr;
                    if(grdftr == _root) _root = nn;
                    else if(grdftr->ftr->lc == grdftr) grdftr->ftr->lc = nn;
                    else grdftr->ftr->rc = nn;
                    connect34(nn, grdftr, pftr, uncle, nn->lc, nn->rc, pftr->rc);
                    nn->RBc = false;
                    grdftr->RBc = true;
                } else {
                    pftr->ftr = grdftr->ftr;
                    if(grdftr == _root) _root = pftr;
                    else if(grdftr->ftr->lc == grdftr) grdftr->ftr->lc = pftr;
                    else grdftr->ftr->rc = pftr;
                    connect34(pftr, grdftr, nn, uncle, pftr->lc, nn->lc, nn->rc);
                    pftr->RBc = false;
                    grdftr->RBc = true;
                }
            }
            return;
        }
    }
}

template <typename T>
void redblacktree<T>::connect34(    Node* nroot,    Node* nlc,      Node* nrc,
                                    Node* ntree1,   Node* ntree2,   Node* ntree3,   Node* ntree4) {
    nlc->lc = ntree1;
    if(ntree1 != NULL) ntree1->ftr = nlc;
    nlc->rc = ntree2;
    if(ntree2 != NULL) ntree2->ftr = nlc;
    nrc->lc = ntree3;
    if(ntree3 != NULL) ntree3->ftr = nrc;
    nrc->rc = ntree4;
    if(ntree4 != NULL) ntree4->ftr = nrc;
    nroot->lc = nlc;
    nlc->ftr = nroot;
    nroot->rc = nrc;
    nrc->ftr = nroot;
    nlc->maintain();
    nrc->maintain();
    nroot->maintain();
}

template <typename T>
typename
redblacktree<T>::iterator redblacktree<T>::lower_bound(T v) {
    Node* ptn = _root;
    while(ptn) {
        _hot = ptn;
        if(ptn->val < v) {
            ptn = ptn->rc;
        } else {
            ptn = ptn->lc;
        }
    }
    if(_hot->val < v) {
        ptn = _hot;
    } else {
        ptn = _hot->left_node();
    }
    return iterator(ptn);
}

template <typename T>
typename
redblacktree<T>::iterator redblacktree<T>::upper_bound(T v) {
    Node* ptn = _root;
    while(ptn) {
        _hot = ptn;
        if(ptn->val > v) {
            ptn = ptn->lc;
        } else {
            ptn = ptn->rc;
        }
    }
    if(_hot->val > v) {
        ptn = _hot;
    } else {
        ptn = _hot->right_node();
    }
    return iterator(ptn);
}

template <typename T>
typename
redblacktree<T>::iterator redblacktree<T>::kth(int rank) {
    return iterator(findkth(rank, _root));
}

template <typename T>
typename
redblacktree<T>::Node* redblacktree<T>::findkth(int rank, Node* ptn) {
    if(!(ptn->lc)) {
        if(rank == 1) {
            return ptn;
        } else {
            return findkth(rank - 1, ptn->rc);
        }
    } else {
        if(ptn->lc->s == rank - 1) return ptn;
        else if(ptn->lc->s >= rank) return findkth(rank, ptn->lc);
        else return findkth(rank - (ptn->lc->s) - 1, ptn->rc);
    }
}

template <typename T>
int redblacktree<T>::get_rank(T v) {
    return find_rank(v, _root);
}

template <typename T>
int redblacktree<T>::find_rank(T v, Node* ptn) {
    if(!ptn) return 1;
    else if(ptn->val >= v) return find_rank(v, ptn->lc);
    else return (1 + ((ptn->lc) ? (ptn->lc->s) : 0) + find_rank(v, ptn->rc));
}

template <typename T>
int redblacktree<T>::size() {
    return _root->s;
}

template <typename T>
bool redblacktree<T>::empty() {
    return !_root;
}

template <typename T>
bool redblacktree<T>::remove(T v) {
    Node* ptn = rfind(v, -1);
    if(!ptn) return false;
    Node* node_suc;
    while(ptn->lc || ptn->rc) { 迭代寻找真后继
        if(!(ptn->lc)) {
            node_suc = ptn->rc;
        } else if(!(ptn->rc)) {
            node_suc = ptn->lc;
        } else {
            node_suc = ptn->succ();
        }
        --(ptn->s);
        ptn->val = node_suc->val;
        ptn = node_suc;
    }
    if(!(ptn->RBc)) {
        --(ptn->s);
        SolveDoubleBlack(ptn);
    }
    if(ptn == _root) {
        _root = NULL;
        delete ptn;
        return true;
    }
    if(ptn->ftr->lc == ptn)
        ptn->ftr->lc = NULL;
    else
        ptn->ftr->rc = NULL;
    delete ptn;
    return true;
}

template <typename T>
void redblacktree<T>::SolveDoubleBlack(Node* nn) {
    while(nn != _root) {
        Node* pftr = nn->ftr;
        Node* bthr = bro(nn);
        if(bthr->RBc) {                 BB-1
            bthr->RBc = false;
            pftr->RBc = true;
            if(_root == pftr) _root = bthr;
            if(pftr->ftr) {
                if(pftr->ftr->lc == pftr)
                    pftr->ftr->lc = bthr;
                else
                    pftr->ftr->rc = bthr;
            }
            bthr->ftr = pftr->ftr;
            if(islc(nn)) {
                connect34(bthr, pftr, bthr->rc, nn, bthr->lc, bthr->rc->lc, bthr->rc->rc);
            } else {
                connect34(bthr, bthr->lc, pftr, bthr->lc->lc, bthr->lc->rc, bthr->rc, nn);
            }
            bthr = bro(nn);
            pftr = nn->ftr;
        }
        if(bthr->lc && bthr->lc->RBc) { BB-3
            bool oldRBc = pftr->RBc;
            pftr->RBc = false;
            if(pftr->lc == nn) {
                if(pftr->ftr) {
                    if(pftr->ftr->lc == pftr)
                        pftr->ftr->lc = bthr->lc;
                    else
                        pftr->ftr->rc = bthr->lc;
                }
                bthr->lc->ftr = pftr->ftr;
                if(_root == pftr) _root = bthr->lc;
                connect34(bthr->lc, pftr, bthr, pftr->lc, bthr->lc->lc, bthr->lc->rc, bthr->rc);
            } else {
                bthr->lc->RBc = false;
                if(pftr->ftr) {
                    if(pftr->ftr->lc == pftr)
                        pftr->ftr->lc = bthr;
                    else
                        pftr->ftr->rc = bthr;
                }
                bthr->ftr = pftr->ftr;
                if(_root == pftr) _root = bthr;
                connect34(bthr, bthr->lc, pftr, bthr->lc->lc, bthr->lc->rc, bthr->rc, pftr->rc);
            }
            pftr->ftr->RBc = oldRBc;
            return;
        } else if(bthr->rc && bthr->rc->RBc) {  BB-3
            bool oldRBc = pftr->RBc;
            pftr->RBc = false;
            if(pftr->lc == nn) {
                bthr->rc->RBc = false;
                if(pftr->ftr) {
                    if(pftr->ftr->lc == pftr)
                        pftr->ftr->lc = bthr;
                    else
                        pftr->ftr->rc = bthr;
                }
                bthr->ftr = pftr->ftr;
                if(_root == pftr) _root = bthr;
                connect34(bthr, pftr, bthr->rc, pftr->lc, bthr->lc, bthr->rc->lc, bthr->rc->rc);
            } else {
                if(pftr->ftr) {
                    if(pftr->ftr->lc == pftr)
                        pftr->ftr->lc = bthr->rc;
                    else
                        pftr->ftr->rc = bthr->rc;
                }
                bthr->rc->ftr = pftr->ftr;
                if(_root == pftr) _root = bthr->rc;
                connect34(bthr->rc, bthr, pftr, bthr->lc, bthr->rc->lc, bthr->rc->rc, pftr->rc);
            }
            pftr->ftr->RBc = oldRBc;
            return;
        }
        if(pftr->RBc) {                 BB-2R
            pftr->RBc = false;
            bthr->RBc = true;
            return;
        } else {                        BB-2B
            bthr->RBc = true;
            nn = pftr;
        }
    }
#ifdef __REDBLACK_DEBUG
    --blackheight;
#endif
}

#ifdef __REDBLACK_DEBUG

int blackheight(0);

template <typename T>   先序遍历
void redblacktree<T>::previs(Node* ptn, int cnt) {
    if(ptn == NULL) {
        if(blackheight == -1) blackheight = cnt;
        assert(blackheight == cnt);
        return;
    }
    printf("%d %s %d \n", ptn->val, ptn->RBc ? "Red" : "Black", ptn->s);
    if(!(ptn->RBc)) ++cnt;
    previs(ptn->lc, cnt);
    previs(ptn->rc, cnt);
}

template <typename T>   中序遍历
void redblacktree<T>::invis(Node* ptn, int cnt) {
    if(ptn == NULL) {
        if(blackheight == -1) blackheight = cnt;
        assert(blackheight == cnt);
        return;
    }
    if(!(ptn->RBc)) ++cnt;
    invis(ptn->lc, cnt);
    printf("%d %s %d \n", ptn->val, ptn->RBc ? "Red" : "Black", ptn->s);
    invis(ptn->rc, cnt);
}

template <typename T>   后序遍历
void redblacktree<T>::postvis(Node* ptn, int cnt) {
    if(ptn == NULL) {
        if(blackheight == -1) blackheight = cnt;
        assert(blackheight == cnt);
        return;
    }
    if(!(ptn->RBc)) ++cnt;
    postvis(ptn->lc, cnt);
    postvis(ptn->rc, cnt);
    printf("%d %s %d \n", ptn->val, ptn->RBc ? "Red" : "Black", ptn->s);
}

template <typename T>   输出所有序遍历的接口
void redblacktree<T>::vis() {
    printf("BlackHeight:\t%d\n", blackheight);
    printf("------pre-vis------\n");
    previs(_root, 0);
    printf("------in-vis------\n");
    invis(_root, 0);
    printf("------post-vis------\n");
    postvis(_root, 0);
}

template <typename T>   验证所有节点与父亲的连接是否正常、域s是否维护正常
void redblacktree<T>::checkconnect(Node* ptn) {
    if(!ptn) return;
    assert(ptn->s > 0);
    if(ptn->lc && ptn->lc->ftr != ptn) {
        printf("Oops! %d has a lc %d, but it failed to point its ftr!\n", ptn->val, ptn->lc->val);
    }
    if(ptn->rc && ptn->rc->ftr != ptn) {
        printf("Oops! %d has a rc %d, but it failed to point its ftr!\n", ptn->val, ptn->rc->val);
    }
    int sss = ptn->s;
    if(ptn->lc) sss -= ptn->lc->s;
    if(ptn->rc) sss -= ptn->rc->s;
    if(sss - 1) {
        printf("Damn! %d's size is %d, but the sum of its children's size is %d!\n", ptn->val, ptn->s, ptn->s - sss);
    }
    checkconnect(ptn->lc);
    checkconnect(ptn->rc);
}

template <typename T>
void redblacktree<T>::correctlyconnected() {
    checkconnect(_root);
}
#endif

#include<iostream>
const int maxn=2e5+10;
int n;long long t,a[maxn],sum[maxn];
redblacktree<long long> tree;


int main()
{
	scanf("%d %lld",&n,&t); long long ans=0;
	tree.insert(0);
	for(int i=1;i<=n;i++){
		scanf("%lld",&a[i]);sum[i]=sum[i-1]+a[i];

		ans+=tree.get_rank(t-sum[i])-1LL;
		//cout << ans << endl;
		tree.insert(-sum[i]);
	}
	printf("%lld\n",ans);
}