【Leetcode】1093. Statistics from a Large Sample

题目地址：

https://leetcode.com/problems/statistics-from-a-large-sample/

给定一个描述某个样本的长$n$数组$A$，$A[i]$代表$i$出现的次数。问这个样本的最小值、最大值、平均值、中位数和众数。题目保证众数一定存在。返回一个浮点型数组。

最小值、最大值和众数很容易求。要求平均值，要先求总和${\sum }_{i=0}^{n-1}iA[i]$和总个数${\sum }_{i=0}^{n-1}A[i]$，则平均值就是$\frac{{\sum }_{i=0}^{n-1}iA[i]}{{\sum }_{i=0}^{n-1}A[i]}$。要求中位数，则需要考虑总个数的奇偶性，如果总个数$k={\sum }_{i=0}^{n-1}A[i]$是奇数，则$A$前缀和第一次超过$k/2+1$的位置即为所求；如果$k$是偶数，则$A$前缀和第一次超过$k/2$的位置和第一次超过$k/2+1$的位置两者平均值即为所求。代码如下：

class Solution {
 public:
  vector<double> sampleStats(vector<int>& cnt) {
    int m = INT_MAX, M = INT_MIN, mode = 0, amount = 0, max_cnt = 0;
    double sum = 0.0;
    for (int i = 0; i < cnt.size(); i++) {
      if (cnt[i]) {
        m = min(m, i), M = max(M, i);
        sum += (double)i * cnt[i];
        amount += cnt[i];
        if (cnt[i] > max_cnt) {
          max_cnt = cnt[i];
          mode = i;
        }
      }
    }

    double ave = sum / amount;
    // 返回序列里第k个数
    auto f = [&](int k) {
      int amount = 0;
      for (int i = m; i <= M; i++) {
        amount += cnt[i];
        if (amount >= k) return i;
      }

      return -1;
    };

    double mean = 0.0;
    if (amount % 2) mean = f(amount / 2 + 1);
    else mean = (f(amount / 2) + f(amount / 2 + 1)) / 2.0;

    return {(double)m, (double)M, ave, mean, (double)mode};
  }
};

时间复杂度$O(n)$，空间$O(1)$。