P160 0 到 100 的数据流中查找中位数

描述#

这个题目说的是，你要设计一个数据结构用于查找数据流中的中位数。它要支持以下两个操作，addNum 用于向数据结构中添加一个整数。findMedian 用于查找该数据结构中目前所有整数的中位数。

注意，这个题目要求你在两种情况下设计这个数据结构。第一种情况是，数据流中的数字全部都大于等于 0 并且小于等于 100。第二种情况是，数据流中 99% 的数字大于等于 0 并且小于等于 100。

其中，中位数的定义是将序列排序后正中间的数字，如果序列的长度是偶数，则取正中间两个数字的平均数。

代码#

public class AlgoCasts {

  public class MedianFinderBetweenZeroAndOneHundred {

    private final int LENGTH = 101;
    private final int[] dataCount = new int[LENGTH];
    private int n = 0;

    // k starts from 1
    // suppose that k is always valid
    // Time: O(1)
    private int findKth(int k) {
      int cnt = 0;
      for (int i = 0; i < LENGTH; ++i) {
        cnt += dataCount[i];
        if (cnt >= k) return i;
      }
      return -1;
    }

    // Time: O(1)
    public void addNum(int num) {
      ++dataCount[num];
      ++n;
    }

    // Time: O(1)
    public double findMedian() {
      if ((n & 1) == 1) return findKth(n/2 + 1);
      else return (findKth(n/2) + findKth(n/2 + 1)) / 2.0;
    }
  }

  public class MedianFinderBetweenZeroAndOneHundred99Percent {

    class MedianFinderArray {
      private final List<Integer> data = new ArrayList<>();

      public void addNum(int num) {
        int idx = data.size() - 1;
        while (idx >= 0 && data.get(idx) > num) --idx;
        data.add(idx+1, num);
      }

      // k starts from 1
      public int findKth(int k) {
        return data.get(k-1);
      }

      public int size() {
        return data.size();
      }
    }

    class MedianFinderBetweenZeroAndOneHundred {
      private final int LENGTH = 101;
      private final int[] dataCount = new int[LENGTH];
      private int n = 0;

      // k starts from 1
      public int findKth(int k) {
        int cnt = 0;
        for (int i = 0; i < LENGTH; ++i) {
          cnt += dataCount[i];
          if (cnt >= k) return i;
        }
        return -1;
      }

      public void addNum(int num) {
        ++dataCount[num];
        ++n;
      }

      public int size() {
        return n;
      }
    }

    private final MedianFinderArray leftData = new MedianFinderArray();
    private final MedianFinderArray rightData = new MedianFinderArray();
    private final MedianFinderBetweenZeroAndOneHundred middleData = new MedianFinderBetweenZeroAndOneHundred();

    // Time(avg): O(1), Time(worst): O(n)
    public void addNum(int num) {
      if (num >= 0 && num <= 100) middleData.addNum(num);
      else if (num < 0) leftData.addNum(num);
      else rightData.addNum(num);
    }

    // Time: O(1)
    public double findMedian() {
      int middleSize = middleData.size();
      int leftSize = leftData.size();
      int rightSize = rightData.size();
      int n = middleSize + leftSize + rightSize;
      if ((n & 1) == 1) {
        int k = n/2 + 1;
        if (k > leftSize && k <= leftSize+middleSize) {
          return middleData.findKth(k - leftSize);
        } else if (k <= leftSize) {
          return leftData.findKth(k);
        } else {
          return rightData.findKth(k - leftSize - middleSize);
        }
      } else {
        int k = n/2;
        if (k > leftSize && k+1 <= leftSize+middleSize) {
          k = k - leftSize;
          return (middleData.findKth(k) + middleData.findKth(k+1)) / 2.0;
        } else if (k+1 <= leftSize) {
          return (leftData.findKth(k) + leftData.findKth(k+1)) / 2.0;
        } else if (k == leftSize && middleSize > 0) {
          return (leftData.findKth(k) + middleData.findKth(1)) / 2.0;
        } else if (k == leftSize && middleSize == 0) {
          return (leftData.findKth(k) + rightData.findKth(1)) / 2.0;
        } else if (k == leftSize + middleSize) {
          return (middleData.findKth(k-leftSize) + rightData.findKth(1)) / 2.0;
        } else {
          k = k - leftSize - middleSize;
          return (rightData.findKth(k) + rightData.findKth(k+1)) / 2.0;
        }
      }
    }

  }

}