/** * Copyright (C) 2014-2016 LinkedIn Corp. (pinot-core@linkedin.com) * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package com.linkedin.pinot.core.util; import com.linkedin.pinot.common.utils.Pairs.IntDoublePair; import it.unimi.dsi.fastutil.doubles.DoubleArrayList; import java.util.Collections; import javax.annotation.concurrent.NotThreadSafe; /** * Heap based Indexed priority queue with primitive 'int' key and 'double' value. * * Allows for the following: * <ul> * <li> O(1) access to values inserted using their corresponding keys. </li> * <li> Dynamic update of values (heap order is maintained after updates). </li> * <li> Min or max ordering, can be specified in the constructor. </li> * </ul> */ @NotThreadSafe @SuppressWarnings("Duplicates") public class IntDoubleIndexedPriorityQueue extends BaseIndexedPriorityQueue { DoubleArrayList _values; IntDoublePair _reusablePair; /** * Constructor for the class. * * @param initialCapacity Initial capacity for the priority queue * @param minHeap Min order, ie smallest element on top. */ public IntDoubleIndexedPriorityQueue(int initialCapacity, boolean minHeap) { super(initialCapacity, minHeap); _values = new DoubleArrayList(initialCapacity); _reusablePair = new IntDoublePair(0, 0.0); } /** * Puts the element into the priority queue. * <ul> * <li> If key does not exist, it is added to the priority queue. </li> * <li> If key exists, then the value is updated, and the priority queue ordering is maintained. </li> * <li> Runtime complexity of {@code O(log(n)}). </li> * </ul> * @param key Integer key for the value * @param value Double value of the key */ public void put(int key, double value) { if (!_keyToIndexMap.containsKey(key)) { _values.add(value); int last = _values.size() - 1; updateKeyIndexMap(key, last); siftUp(last); } else { int index = _keyToIndexMap.get(key); _values.set(index, value); // Sift the value up or down, as the case may be. if (!siftDown(index)) { siftUp(index); } } } /** * Returns the value for the specified key. * <ul> * <li> Returns null if the specified key does not exist. </li> * <li> Runtime complexity of O(1). </li> * </ul> * * @param key Key for which to return the value * @return Value for the key */ public IntDoublePair get(int key) { if (!_keyToIndexMap.containsKey(key)) { return null; } int index = _keyToIndexMap.get(key); double value = _values.getDouble(index); _reusablePair.setIntValue(index); _reusablePair.setDoubleValue(value); return _reusablePair; } /** * Returns the key+value pair with the max priority (min for minHeap mode) * <ul> * <li> key+value pair is removed from the priority queue. </li> * <li> Returns null if the priority queue is empty. </li> * <li> Runtime complexity of O(1). </li> * </ul> * * @return Key+Value pair */ public IntDoublePair poll() { if (isEmpty()) { return null; } IntDoublePair poll = peek(); int lastIndex = _values.size() - 1; swapValues(0, lastIndex); _values.removeDouble(lastIndex); _keyToIndexMap.remove(_indexToKeyMap.get(lastIndex)); _indexToKeyMap.remove(lastIndex); if (!_values.isEmpty()) { siftDown(0); } return poll; } /** * Returns the key+value pair with the max priority (min for minHeap mode) * <ul> * <li> key+value pair is not removed from the priority queue. </li> * <li> Throws runtime exception if the priority queue is empty. </li> * <li> Runtime complexity of O(1). </li> * </ul> * * @return Key+Value pair */ public IntDoublePair peek() { if (_values.isEmpty()) { throw new RuntimeException("Empty collection"); } _reusablePair.setIntValue(_indexToKeyMap.get(0)); _reusablePair.setDoubleValue(_values.getDouble(0)); return _reusablePair; } /** * Returns true if the priority queue is empty, false otherwise. * * @return True if empty, false otherwise */ public boolean isEmpty() { return _values.isEmpty(); } /** * Helper method that moves the element at the specified index up * until the heap ordering is established. * * @param index Index of element to sift up. */ private void siftUp(int index) { // Return if already at root node. if (index == 0) { return; } while (index != 0) { int parentIndex = getParentIndex(index); double value = _values.getDouble(index); double parentValue = _values.getDouble(parentIndex); if (compare(parentValue, value) == 1) { swapValues(index, parentIndex); index = parentIndex; } else { // No more sifting up required, break break; } } } /** * Helper method that moves the element at the specified index down * until the heap ordering is established. * * @param index Index of element to sift down. * @return True if sifted, false otherwise. */ private boolean siftDown(int index) { boolean hasChildren = hasChildren(index); if (!hasChildren) { return false; } boolean sifted = false; while (true) { int leftChildIndex = getLeftChildIndex(index); int rightChildIndex = getRightChildIndex(index); int minIndex; int size = _values.size(); if (leftChildIndex >= size && rightChildIndex >= size) { // This is leaf node, all done. break; } else if (rightChildIndex >= size) { // Node only has left child which will be the minimum. minIndex = leftChildIndex; } else { // Node has both left and right children, find the minimum of the two. double leftChildValue = _values.getDouble(leftChildIndex); double rightChildValue = _values.getDouble(rightChildIndex); if (compare(leftChildValue, rightChildValue) <= 0) { minIndex = leftChildIndex; } else { minIndex = rightChildIndex; } } // One of the children is out of order, need to sift down. if (compare(_values.getDouble(index), _values.getDouble(minIndex)) == 1) { swapValues(index, minIndex); index = minIndex; sifted = true; } else { break; } } return sifted; } /** * Compares the two specified values, and returns: * <ul> * <li> if v1 < v2, -1 for max, +1 for min mode. </li> * <li> if v1 > v2, -1 for max, -1 for min mode. </li> * <li> if v1 = v2, 0 for max, 0 for min mode. </li> * </ul> * @param v1 Value to compare * @param v2 Value to compare * @return Result of comparison (as described above). */ private int compare(double v1, double v2) { int ret = Double.compare(v1, v2); return (_minHeap) ? ret : -ret; } /** * Helper method that performs all operations required to swap two values. * <ul> * <li> Swaps the values in the array that backs the heap. </li> * <li> Updates the indexToKey and keyToIndex maps due to the swap. </li> * </ul> * @param index1 Index to swap * @param index2 Index to swap */ private void swapValues(int index1, int index2) { if (index1 == index2) { return; } double tmp = _values.getDouble(index1); _values.set(index1, _values.getDouble(index2)); _values.set(index2, tmp); swapKeys(index1, index2); } /** * Returns true if the node at specified index has children, false otherwise. * Just checking for existence of left child is sufficient (array backed heap). * * @param index Index to check * @return True if node has children, false otherwise. */ private boolean hasChildren(int index) { return (getLeftChildIndex(index) < _values.size()); } }