3369. Design an Array Statistics Tracker

• Time:
  • Constructor: $O(n)$
  • addNumber(number: int): $O(\log \texttt{addNumber()})$
  • removeFirstAddedNumber(): $O(\log \texttt{addNumber(number: int)})$
  • getMean(): $O(1)$
  • getMedian(): $O(1)$
  • getMode(): $O(\texttt{addNumber(number: int)})$
• Space: $O(\texttt{addNumber(number: int)})$

C++JavaPython

class StatisticsTracker {
 public:
  void addNumber(int number) {
    q.push(number);
    ++count[number];
    sortedList.insert(number);
    modeMaxHeap.emplace(count[number], -number);
    sum += number;
  }

  void removeFirstAddedNumber() {
    const int number = q.front();
    q.pop();
    if (--count[number])
      count.erase(number);
    sortedList.erase(sortedList.find(number));
    // Note: No need to update the heap now; we'll clean up stale entries when
    // getting the mode.
    sum -= number;
  }

  int getMean() {
    return sum / q.size();
  }

  int getMedian() {
    auto it = sortedList.begin();
    advance(it, sortedList.size() / 2);
    return *it;
  }

  int getMode() {
    // Removes stale entries from the top of the heap.
    while (!modeMaxHeap.empty()) {
      const int frequency = modeMaxHeap.top().first;
      const int number = -modeMaxHeap.top().second;
      if (count.contains(number) && count[number] == frequency)
        return number;
      modeMaxHeap.pop();
    }
    throw;
  }

 private:
  queue<int> q;
  unordered_map<int, int> count;
  multiset<int> sortedList;
  priority_queue<pair<int, int>> modeMaxHeap;  // (frequency, number)
  long sum = 0;
};

class StatisticsTracker {
  public void addNumber(int number) {
    q.offer(number);
    count.merge(number, 1, Integer::sum);
    sortedList.merge(number, 1, Integer::sum);
    modeMaxHeap.offer(new Pair<>(count.get(number), number));
    sum += number;
  }

  public void removeFirstAddedNumber() {
    final int number = q.poll();
    count.merge(number, -1, Integer::sum);
    if (count.get(number) == 0)
      count.remove(number);
    sortedList.merge(number, -1, Integer::sum);
    if (sortedList.get(number) == 0)
      sortedList.remove(number);
    sum -= number;
  }

  public int getMean() {
    return (int) (sum / q.size());
  }

  public int getMedian() {
    final int midIndex = q.size() / 2; // Median depends on the queue size.
    int count = 0;
    for (Map.Entry<Integer, Integer> entry : sortedList.entrySet()) {
      count += entry.getValue();
      if (count > midIndex)
        return entry.getKey();
    }
    throw new IllegalArgumentException();
  }

  public int getMode() {
    // Removes stale entries from the top of the heap.
    while (!modeMaxHeap.isEmpty()) {
      final int frequency = modeMaxHeap.peek().getKey();
      final int number = modeMaxHeap.peek().getValue();
      if (count.containsKey(number) && count.get(number) == frequency)
        return number;
      modeMaxHeap.poll();
    }
    throw new IllegalArgumentException();
  }

  private Queue<Integer> q = new LinkedList<>();
  private Map<Integer, Integer> count = new HashMap<>();
  private TreeMap<Integer, Integer> sortedList = new TreeMap<>();
  // (frequency, number)
  private PriorityQueue<Pair<Integer, Integer>> modeMaxHeap =
      new PriorityQueue<>(Comparator.comparingInt(Pair<Integer, Integer>::getKey)
                              .reversed()
                              .thenComparingInt(Pair<Integer, Integer>::getValue));
  private long sum = 0;
}

from sortedcontainers import SortedList


class StatisticsTracker:
  def __init__(self):
    self.q = collections.deque()
    self.count = collections.Counter()
    self.sortedList = SortedList()
    self.modeMaxHeap = []  # (frequency, number)
    self.sum = 0

  def addNumber(self, number: int) -> None:
    self.q.append(number)
    self.count[number] += 1
    self.sortedList.add(number)
    heapq.heappush(self.modeMaxHeap, (-self.count[number], number))
    self.sum += number

  def removeFirstAddedNumber(self) -> None:
    number = self.q.popleft()
    self.count[number] -= 1
    self.sortedList.remove(number)
    # Note: No need to update the heap now; we'll clean up stale entries when
    # getting the mode.
    self.sum -= number

  def getMean(self) -> int:
    return self.sum // len(self.q)

  def getMedian(self) -> int:
    return self.sortedList[len(self.sortedList) // 2]

  def getMode(self) -> int:
    # Removes stale heap entries where frequency no longer matches.
    while self.modeMaxHeap:
      frequency = -self.modeMaxHeap[0][0]
      number = self.modeMaxHeap[0][1]
      if self.count[number] == frequency:
        return number
      heapq.heappop(self.modeMaxHeap)