Одно из возможных решений — использовать две кучи разных приоритетов: максимальная куча (maxHeap) для значений выше среднего и минимальная куча (minHeap) для значений ниже среднего. Это позволит разделить элементы примерно поровну с двумя значениями — вершинами куч. Теперь найти среднее значение очень просто.
Что означает «примерно поровну»? «Примерно» означает, что при нечетном количестве чисел в одной из куч окажется лишнее число. Можно сформулировать:
если maxHeap.size() > min.Heap.size(), то heap1.top() будет средним значением;
если maxHeap.size() == minHeap.size(), то средним значением будет среднее значений maxHeap.top() и min.Heap.top().
В алгоритме с балансировкой мы гарантируем, что maxHeap будет всегда содержать дополнительный элемент.
Алгоритм работает следующим образом. Если новое значение меньше или равно среднему, оно помещается в maxHeap, в противном случае оно попадает в minHeap. Размеры куч могут совпадать или в maxHeap может быть один дополнительный элемент. Это требование легко выполнить, сдвигая элемент из одной кучи в другую. Среднее значение находится в вершине. Обновления занимают O(log(n)) времени.
private Comparator<Integer> maxHeapComparator;
private Comparator<Integer> minHeapComparator;
private PriorityQueue<Integer> maxHeap, minHeap;
public void addNewNumber(int randomNumber) {
/*Заметьте: addNewNumber поддерживает условие, что
*maxHeap.size() >= minHeap.size() */
if (maxHeap.size() >= minHeap.size()) {
if ((minHeap.peek() != null) &&
randomNumber > minHeap.peek()) {
maxHeap.offer(minHeap.poll());
minHeap.offer(randomNumber);
} else {
maxHeap.offer(randomNumber);
}
} else {
if(randomNumber < maxHeap.peek()) {
minHeap.offer(maxHeap.poll());
maxHeap.offer(randomNumber)
}
else {
minHeap.offer(randomNumber);
}
}
}
public static double getMedian() {
/*maxHeap является всегда по крайней мере столь же большой,
*как minHeap. Если maxHeap пуста, то minHeap тоже пуста. */
if (maxHeap.isEmpty()) {
return 0;
}
if(maxHeap.size() == minHeap.size()) {
return ((double)minHeap.peek()+(double)maxHeap.peek()) / 2;
} else {
/* Если maxHeap и minHeap разных размеров, то
* в maxHeap есть один дополнительный элемент.
* Возвращаем вергину кучи maxHeap */
return maxHeap.peek();
}
}
Разбор задачи по книге «Карьера программиста. Как устроиться на работу в Google, Microsoft или другую ведущую IT-компанию»