【每日一题Day91】LC1825求出 MK 平均值 | 有序集合+模拟

求出 MK 平均值【LC1825】

给你两个整数 m 和 k ，以及数据流形式的若干整数。你需要实现一个数据结构，计算这个数据流的 MK 平均值 。

MK 平均值 按照如下步骤计算：

1. 如果数据流中的整数少于 m 个，MK 平均值 为 -1 ，否则将数据流中最后 m 个元素拷贝到一个独立的容器中。
2. 从这个容器中删除最小的 k 个数和最大的 k 个数。
3. 计算剩余元素的平均值，并 向下取整到最近的整数 。

请你实现 MKAverage 类：

• MKAverage(int m, int k) 用一个空的数据流和两个整数 m 和 k 初始化 MKAverage 对象。
• void addElement(int num) 往数据流中插入一个新的元素 num 。
• int calculateMKAverage() 对当前的数据流计算并返回 MK 平均数 ，结果需 向下取整到最近的整数 。

重新投入学习

数组+排序

• 思路：使用一个长度为$m$的数组记录数据流的最近$m$位数据，然后每次计算平均值时，将数据排序，然后计算中间$[k+1,m-k-1]$个数据之和，然后计算平均值

• 实现【超时】

  $i$为总共插入的元素下标，通过取余的方式覆盖之前的元素

  class MKAverage {
      int[] nums;
      int[] sort;
      int i, m, k;
      public MKAverage(int m, int k) {
          nums = new int[m];
          sort = new int[m];
          i = 0;
          this.m = m;
          this.k = k;
      }
      
      public void addElement(int num) {
          nums[i % m] = num;
          i++;
      }
      
      public int calculateMKAverage() {
          if (i < m) return -1;
          for (int i = 0; i < m; i++){
              sort[i] = nums[i];
          }
          Arrays.sort(sort);       
          int sum = 0;
          for (int i = k; i < m - k; i++){
              sum += sort[k];
          }
          
          return sum / (m - 2 * k);
      }
  }
  
  /**
   * Your MKAverage object will be instantiated and called as such:
   * MKAverage obj = new MKAverage(m, k);
   * obj.addElement(num);
   * int param_2 = obj.calculateMKAverage();
   */
  

  • 复杂度

    • 时间复杂度：addElement函数的时间复杂度为$O(1)$。calculateMKAverage函数的时间复杂度为$O(mlogm)$。
    • 空间复杂度：$O(m)$

TreeMap

• 思路：排序后的数组可以分为三个部分：较小的前k个元素、中间元素、前大的后k个元素，因此使用3个有序集合存储这三个部分。在加入元素时，判断元素的个数以及元素的大小移动元素，时间复杂度为$O(logm)$；在求平均值时，直接求取中间元素的平均值，时间复杂度为$O(1)$。

• 实现

  维护一下数据结构和变量

  • 使用三个有序集合TreeMapmin,mid,max分别存储最小的k个元素、中间的元素、最大的k个元素，key值为元素大小，value值为元素个数，因此还需要使用额外变量记录集合中元素的个数
  • 使用一个队列记录当前的m个元素，队列头部为最早加入的元素，队尾元素为最近加入的元素
  • 使用一个int类型变量记录中间元素之和

  函数实现

  • addElement

    1. 加入元素num值对应位置，并放入队列中：
       • 如果min为空，或者$num\le max(min)$，则将num加入min中
       • 否则，如果max为空，或者$num\ge min(max)$，则将num加入max中
       • 否则，则将num加入mid中，sum增加num
    2. 移除元素：
       • 如果队列的长度大于m，那么移除队首元素，查找队首元素位于哪个集合中，并从该集合中删除该元素，如果该集合为mid，则sum减去队首元素
       • 并判断三个集合中的元素个数是否大于或者小于k，若大于k，那么先移动元素至mid，再将元素移动至min或者max
         • 如果min的长度大于k，则将最大值放入mid中
         • 如果max的长度大于k，则将最小值放入mid中
         • 如果min的长度小于k，则将mid的最小值放入min中
         • 如果max的长度小于k，则将mid的最大值放入max中

  • calculateMKAverage

    如果队列长度小于m，那么返回-1，否则返回$sum/(m-2\ast k)$

• 实现

  class MKAverage {
      TreeMap<Integer,Integer> min = new TreeMap<>();
      TreeMap<Integer,Integer> mid = new TreeMap<>();
      TreeMap<Integer,Integer> max = new TreeMap<>();
      Deque<Integer> dq = new LinkedList<>();
      int sum = 0, s1 = 0, s2 = 0, s3 = 0;
      int m, k;
      public MKAverage(int m, int k) {
          this.m = m;
          this.k = k;
      }
      
      public void addElement(int num) {
          // 放入有序集合中
          if (s1 == 0 || num <= min.lastKey()){
              min.merge(num, 1, Integer::sum);
              s1++;
          }else if (s3 == 0 || num >= max.firstKey()){
              max.merge(num, 1, Integer::sum);
              s3++;
          }else{
              mid.merge(num, 1, Integer::sum);
              sum += num;
              s2++;
          }
          // 放入队列中
          dq.addLast(num);
          // 删除队首元素
          if (dq.size() > m){
              int poll = dq.pollFirst();
              if (min.containsKey(poll)){
                  if (min.merge(poll, -1, Integer::sum) == 0){
                      min.remove(poll);
                  }
                  s1--;
              }else if (max.containsKey(poll)){
                  if (max.merge(poll, -1, Integer::sum) == 0){
                      max.remove(poll);
                  }
                  s3--;
              }else{
                  if (mid.merge(poll, -1, Integer::sum) == 0){
                      mid.remove(poll);
                  }
                  s2--;
                  sum -= poll;
              }
          }
          // 调整有序集合中的元素
          while (s1 > k){
              int x = min.lastKey();
              if (min.merge(x, -1 ,Integer::sum) == 0){
                  min.remove(x);
              }
              mid.merge(x, 1, Integer::sum);
              sum += x;
              s2++;
              s1--;
          }
          while (s3 > k){
              int x = max.firstKey();
              if (max.merge(x, -1 ,Integer::sum) == 0){
                  max.remove(x);
              }
              mid.merge(x, 1, Integer::sum);
              sum += x;
              s2++;
              s3--;
          }
          while (s1 < k && !mid.isEmpty()){
              int x = mid.firstKey();
              if (mid.merge(x, -1 ,Integer::sum) == 0){
                  mid.remove(x);
              }
              sum -= x;
              min.merge(x, 1, Integer::sum);
              s2--;           
              s1++;
          }
          while (s3 < k && !mid.isEmpty()){
              int x = mid.lastKey();
              if (mid.merge(x, -1 ,Integer::sum) == 0){
                  mid.remove(x);
              }
              sum -= x;
              max.merge(x, 1, Integer::sum);
              s2--;           
              s3++;
          }
      }
      
      public int calculateMKAverage() {
          if (dq.size() < m){
              return -1;
          }
          return sum / (m - 2 * k);
          
      }
  }
  
  /**
   * Your MKAverage object will be instantiated and called as such:
   * MKAverage obj = new MKAverage(m, k);
   * obj.addElement(num);
   * int param_2 = obj.calculateMKAverage();
   */