牛客87-B

K小数查询

---

题目描述

给出长度为 n 的序列 $a_1$, $a_2$, ..., $a_n$,有多少对整数(l,r),满足 r-l+1>=k且 $a_l$, $a_{l+1}$, ...,$a_r$中第 k 小的数是 x？
输入描述

• 第一行三个整数 n,x,k
• 第二行为$a_1$, $a_2$, ..., $a_n$
  1 <= x,k <= n <= 200000,
  $a_1$, $a_2$, ..., $a_n$为1，2，3...，n的一个排列
  原题链接

---

Input

5 3 2
1 2 3 4 5

10 5 2
8 9 10 2 4 5 1 6 3 7

Output

3
3

解题思路

题目意思就是给出一个区间[l,r],让 x 为这个区间的第 k 小，求这样的区间有多少个。既然 x 是第 k 小，那么区间[l,r]一定要包含 x ，则可以分为三种情况。假设 x 是给出序列的第 dis 个元素(即$a_{dis}$ = x)。

• l = dis时，有多少个 r 使得 $a_{dis+1}$,$a_{dis+2}$,...$a_r$中,有 k - 1 个元素小于 x
• 当 r = dis时，有多少个 l 使得 $a_l$,$a_{l+1}$,...$a_{dis-1}$中，有 k - 1个元素小于x
• 当 l < dis,r > dis 时，有多少对[l,r]使得，$a_l$, $a_{l+1}$,..., $a_{dis-1}$中，有 a 个元素小于x，$a_{dis+1}$, $a_{dis+2}$,..., $a_r$中,有 b 个元素小于 x
  且 a + b == k - 1
  我们可以开一个数组sum[N]， sum[i]表示 [i,dis)（或(dis,i]),中有多少个元素小于 x
  如图 n = 10,x = 5, k = 2
  
  直接用sum数组进行计算的话会超时，所以我又开了两个数组 l_vis[],r_vis[]。 l_vis[i]储存 [1,dis),sum = [i]的个数，r_vis[i]储存(dis,n],sum = i的个数。ans += l_vis[i] * r_vis[k-1-i], i 从 0枚举到 k - 1即可，枚举完毕后，这个只是计算完了第三种情况，ans += l_vis[k-1] + r_vis[k-1],这样 1，2情况也包含在内了

AC代码

#include <iostream>
using namespace std;
const int N = 2e5 + 10;
int arr[N],sum[N],n,x,k,dis;
int l_vis[N],r_vis[N];
int main(){
    long long ans = 0;
    scanf("%d %d %d",&n,&x,&k);
    for(int i = 1; i <= n; i++){
        scanf("%d",&arr[i]);
        if(arr[i] == x) dis = i;
    }
    for(int i = dis - 1; i >= 1; i--) sum[i] = sum[i+1] + (arr[i] < x);
    for(int i = dis + 1; i <= n; i++) sum[i] = sum[i-1] + (arr[i] < x);
    for(int i = 1; i < dis; i++) l_vis[sum[i]] ++;
    for(int i = dis + 1; i <= n; i++) r_vis[sum[i]]++;
    for(int i = 0; i < k; i++)
        ans += l_vis[i] * r_vis[k-i-1];
    ans += l_vis[k-1] + r_vis[k-1];
    printf("%lld",ans);
    return 0;
}