Meaningful Mean AtCoder - 2581 （樹狀數組+離散化）

Problem Statement

You are given an integer sequence of length N, a= {a₁,a₂,…,a_N}, and an integer K.

a has N(N+1)⁄2 non-empty contiguous subsequences, {a_l,a_l+1,…,a_r} (1≤l≤r≤N). Among them, how many have an arithmetic mean that is greater than or equal to K?

Constraints

• All input values are integers.
• 1≤N≤2×10⁵
• 1≤K≤10⁹
• 1≤a_i≤10⁹

Input

Input is given from Standard Input in the following format:

N K
a1
a2
:
aN

Output

Print the number of the non-empty contiguous subsequences with an arithmetic mean that is greater than or equal to K.

Sample Input 1

3 6
7
5
7

Sample Output 1

5

All the non-empty contiguous subsequences of a are listed below:

• {a₁} = {7}
• {a₁,a₂} = {7,5}
• {a₁,a₂,a₃} = {7,5,7}
• {a₂} = {5}
• {a₂,a₃} = {5,7}
• {a₃} = {7}

Their means are 7, 6, 19⁄3, 5, 6 and 7, respectively, and five among them are 6or greater. Note that {a₁} and {a₃} are indistinguishable by the values of their elements, but we count them individually.

Sample Input 2

1 2
1

Sample Output 2

0

Sample Input 3

7 26
10
20
30
40
30
20
10

Sample Output 3

13

題意：給n個數字，現在求 平均數大於K的集合（非空），有多少個？

思路：枚舉所有的區間  （ L,R 】 (左開右閉，爲了使區間不重疊，不同時包含端點） 我們先計算出sum[ i ]數組，代表前i個數的和，那麼（L,R 】區間的和爲sum【R】-sum【L】，sum【R】-sum【L】>=k*(R-L);==>sum【R】- sum【L】>=K*R-K*L；==>sum【R】-K*R>=sum【L】-K*L；(0<=L<R<=n)（L等於0時，第一個數字才能包含在區間裏）(在符合此式子即爲滿足題意的區間）  sum【R】-K*R  (0<=R<=n)可以看成一個整體，先預處理求出所有該式子的值，這裏我用num【i】代表sum【i】-K*i，確定區間右端點R，尋找有多少個左端點L 符合 num【L】<=num【R】，這個用樹狀數組來解決該問題（求R前面有多少個num【L】小於他的num）。因爲num【i】很大，所以先離散化，然後用樹狀數組nlogn處理。

（因爲(0,R],才能表示從a[1]開始的數據，所以num[0]，sum[0]也是有必要算上的）

代碼：

#include <iostream>
#include<stdio.h>
#include<string.h>
#include<math.h>
#include<queue>
#include<algorithm>
#define N 200500
#define LL long long
#define mem(a,b) memset(a,b,sizeof(a));
using namespace std;
LL sum[N],t[N];
LL tree[N];
int n,k,a,m,nn;
void add(int x)
{
    while(x<=nn)//nn爲上限，要大於離散化後的最大的數字，
    {
        tree[x]++;
        x+=x&-x;//向上更新，樹狀數組核心代碼。此處就不解釋了，不會的先學樹狀數組。
    }
}
LL query(int x)
{
    LL num=0;
    while(x)
    {
        num+=tree[x];
        x-=x&-x;//向下求和
    }
    return num;
}
int main()
{
    mem(t,0);
    mem(tree,0);
    mem(sum,0);
    scanf("%d%d",&n,&k);
    nn=n+10;
    for(int i=1; i<=n; i++)
    {
        scanf("%d",&a);
        t[i]=sum[i]=sum[i-1]+a-k;//預處理出sum[i]數組（就是上面講解的num數組）
    }//t[i]數組爲離散化的輔助數組
    sort(t,t+n+1);//排序，排n+1個數字的序（包括了sum[0]）
    m=unique(t,t+n+1)-t;//去重
    for(int i=0; i<=n; i++)
        sum[i]=lower_bound(t,t+m+1,sum[i])-t+1;//離散化
    LL ans=0;
    for(int i=0;i<=n;i++)
    {
        ans+=query(sum[i]);//查詢有多少個小於sum[i]的數字
        add(sum[i]);//更新樹樁數組。
    }
    printf("%lld\n",ans);
}