2488.统计中位数为 K 的子数组 - 哈希表+前缀和

原题链接：https://leetcode.cn/problems/count-subarrays-with-median-k/

题解

基本思路：对于求中位数而言，我们可以先将比k大的数字转换成1，比k小的数字转换成-1；k转换成0，此时当子数组的中位数为k时，可以得出子数组的和为0或1

即[3,2,1,4,5]，k=4有[-1,-1,-1,0,1]

初始解法：i从k开始往左遍历并sum+=nums[i]，每遍历一位j从k开始往右遍历到底，然后sum+=nums[j]，当sum==0或sum==1时cnt++；要注意每次j遍历完要将sum重置以便下次j便利

时间复杂度：O(n^2)

• 每次i往左遍历一位后，j都需要从k往右遍历到底
• 这么做几乎是必然TLE的

  class Solution {
      public int countSubarrays(int[] nums, int k) {
          int kpos=0,i=0,j=0,a=0;
          for(a=0;a<nums.length;a++){
              if(nums[a]>k){
                  nums[a]=1;
              }else if(nums[a]<k){
                  nums[a]=-1;
              }else{
                  nums[a]=0;
                  kpos=a;
              }
          }
          
          /*O(n^2)*/
          int sum=0,cnt=0,temp=0;
          for(i=kpos;i>=0;i--){
              sum=temp+nums[i];
              temp=sum;
              for(j=kpos;j<nums.length;j++){
                  sum+=nums[j];
                  cnt+=getCnt(sum);
              }
              sum=temp;
          }
          
          return cnt;
      }
      
      int getCnt(int sum){
          if(sum==0||sum==1){return 1;}
          else return 0;
      }
  }

哈希表+前缀和：先让i从k往左遍历到头，每遍历一位sum+=nums[i]，记录下每一位的前缀和，然后用哈希表记录下[0,kpos]中所有前缀和的数量；第二次重置sum，让i从k往右遍历到头，每遍历一位sum+=nums[i]，从哈希表中找是否存在键值对满足key+sum=0或1，然后cnt加上键值对的value即可

时间复杂度：O(n)

• 因为i从k往左遍历到底，再从k往右遍历到底，只需要遍历整个数组一次即可

class Solution {
    public int countSubarrays(int[] nums, int k) {
        int kpos=0,i=0,j=0,a=0;
        for(a=0;a<nums.length;a++){
            if(nums[a]>k){
                nums[a]=1;
            }else if(nums[a]<k){
                nums[a]=-1;
            }else{
                nums[a]=0;
                kpos=a;
            }
        }
        
        /*O(n)*/
        Map<Integer,Integer> map=new HashMap<>();
        int sum=0,cnt=0;
        for(i=kpos;i>=0;i--){
            sum+=nums[i];
            if(!map.containsKey(sum)){map.put(sum,0);}
            map.put(sum,map.get(sum)+1);
        }
        sum=0;
        for(i=kpos;i<nums.length;i++){
            sum+=nums[i];
            if(map.containsKey(0-sum)){
                cnt+=map.get(0-sum);
            }
            if(map.containsKey(1-sum)){
                cnt+=map.get(1-sum);
            }
        }
        
        return cnt;
    }
}