涼宮春日的嘆息「二分求K小」

題目描述

給定一個數組，將其所有子區間的和從小到大排序，求第 \(k\) 小的是多少。

輸入格式

第一行兩個數 \(n\), \(k\)，表示陣列的長度和k；

第二行有 \(n\) 個數，第 \(i\) 個是 \(a[i]\)，表示給定的陣列。

輸出格式

僅一個數，表示答案。

樣例

樣例輸入1

5 6
1 1 1 1 1

樣例輸出1

2

樣例輸入2

8 20
2 3 1 2 5 3 2 3

樣例輸出2

8

資料範圍

對於 \(100\%\) 的資料， \(n <= 10^6\),\(1<=a[i]\),\(k <= 10^9\)

思路分析

• 好像前段時間學長在有一次講課的時候說了一下求 \(k\)
  大/小 的方法，其中有一個就是二分
• 所以這題我們考慮怎麼用二分做，檢驗肯定就是列舉 \(mid\) 看能不能湊出k個小於等於 \(mid\) 的即可
• 但是不同的區間是有很多的，如果枚舉出所有再進行二分檢驗，顯然時間效率是不允許的。但其實有一個很顯然的性質，就是如果一個區間符合條件(即比 \(mid\) 要小)，那麼該區間內的所有區間都會滿足條件。
• 用字首和，固定右端點變化坐端點就彳亍了，可以快速得出答案

Code

#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
#define int long long //奇巧淫技
#define N 1000010
using namespace std;
inline int read(){
	int x=0,f=1;
	char ch = getchar();
	while(ch>'9'||ch<'0'){if(ch=='-')f=-1;ch=getchar();}
	while(ch>='0'&&ch<='9'){x=(x<<1)+(x<<3)+(ch^48);ch=getchar();}
	return x*f;
}
int n,k,a[N],sum[N];
int check(int mid) {
    int j = 0;
    int res = 0;
    for (int i = 1; i <= n; i++) {
        while (sum[i] - sum[j] > mid) j++; //不符合就一直加
        res += i - j; //跳出迴圈後，這個區間裡所有右端點為i的都符合條件
    }
    return res;
}

signed main() {
    n = read(),k = read();
    for (int i = 1; i <= n; i++) {
        a[i] = read();
        sum[i] = sum[i - 1]+a[i];
    }
    int l = 0, r = sum[n], mid;
    while (l < r) {
        mid = (l + r) >> 1;
        if (check(mid) >= k) r = mid;
        else l = mid + 1;
    }
    printf("%lld\n", r);
    return 0;
}