1
\$\begingroup\$

I have implemented this heap (priority queue) data structure using java.util.TreeMap and java.util.ArrayDeque. The operations allowed are:

  • insert(final E element, final P priority: inserts element into the heap assigning priority as the priority key; this operation accepts duplicates,
  • extractMinimum(): removes and returns the head element of the chain with the minimum priority.

The cool part of this implementation is that upon extracting the minimum the time complexity is \$\mathcal{O}(\log P)\$, where \$P\$ is the number of distinct priorities in the heap, and not \$\mathcal{O}(\log n)\$, where \$n\$ is the number of elements stored in the heap, as in java.util.PriorityQueue.

My code is as follows:

JavaHeap.java:

package net.coderodde.util;

import java.util.ArrayDeque;
import java.util.Deque;
import java.util.HashMap;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.PriorityQueue;
import java.util.Random;
import java.util.TreeMap;

/**
 * This class implements a simple priority queue glued from different Java 
 * collections.
 * 
 * @author Rodion "rodde" Efremov
 * @version 1.6 (Jun 19, 2016)
 * @param <E> the type of elements being stored.
 * @param <P> the type of the priority keys.
 */
public class JavaHeap<E, P extends Comparable<? super P>> {

    /**
     * Maps each distinct priority key to the list of elements currently 
     * assigned with that very priority.
     */
    private final TreeMap<P, Deque<E>> mapPriorityToElements = new TreeMap<>();

    /**
     * Caches the amount of elements stored in this heap.
     */
    private int size;

    /**
     * If {@code element} is not present in this heap, adds it and assigns 
     * {@code priority} as its priority. Otherwise, updates the priority of the
     * {@code element}.
     * 
     * @param element  the element to add or update.
     * @param priority the new priority of the element.
     */
    public void insert(final E element, final P priority) {
        Objects.requireNonNull(priority, "The priority key is null.");
        ++size;

        final Deque<E> collisionChain = mapPriorityToElements.get(priority);

        if (collisionChain != null) {
            collisionChain.addLast(element);
        } else {
            final Deque<E> newCollisionChain = new ArrayDeque<>();
            newCollisionChain.addLast(element);
            mapPriorityToElements.put(priority, newCollisionChain);
        }
    }

    public E extractMinimum() {
        checkNotEmpty();

        final Map.Entry<P, Deque<E>> entry = mapPriorityToElements.firstEntry();
        final Deque<E> minimumCollisionChain = entry.getValue();
        final E minimum = minimumCollisionChain.removeFirst();

        if (minimumCollisionChain.isEmpty()) {
            mapPriorityToElements.remove(entry.getKey());
        }

        --size;
        return minimum;
    }

    public int size() {
        return size;
    }

    public boolean isEmpty() {
        return size == 0;
    }

    public void clear() {
        size = 0;
        mapPriorityToElements.clear();
    }

    @Override
    public String toString() {
        final StringBuilder sb = new StringBuilder();

        for (final Map.Entry<P, Deque<E>> entry :
                mapPriorityToElements.entrySet()) {
            for (final E element : entry.getValue()) {
                sb.append("[")
                  .append(element)
                  .append(": p = ")
                  .append(entry.getKey())
                  .append("], ");
            }
        }

        return sb.delete(sb.length() - 2, sb.length()).toString();
    }

    public String toDebugString() {
        if (isEmpty()) {
            return "[empty heap]";
        }

        final StringBuilder sb = new StringBuilder();

        for (final Map.Entry<P, Deque<E>> entry : 
                mapPriorityToElements.entrySet()) {
            toDebugString(sb, entry);
        }

        // Delete the last character that is '\n' and return the string:
        return sb.deleteCharAt(sb.length() - 1).toString();
    }

    private void toDebugString(final StringBuilder sb, 
                               final Map.Entry<P, Deque<E>> entry) {
        sb.append(entry.getKey())
          .append(" -> [");

        int i = 0;
        final int chainLength = entry.getValue().size();

        for (final E element : entry.getValue()) {
            sb.append(element);

            if (i < chainLength - 1) {
                sb.append(", ");
            }

            ++i;
        }

        sb.append("]\n");
    }

    private void checkNotEmpty() {
        if (size == 0) {
            throw new NoSuchElementException("The heap is currently empty.");
        }
    }

    private static final int WARMUP_SIZE    = 1_000_000;
    private static final int BENCHMARK_SIZE = 5_000_000;
    private static final int UNIQUE_VALUES  = 5_000;

    public static void main(String[] args) {
        funnyDemo();
        warmup();
        benchmark();
    }

    private static void funnyDemo() {
        System.out.println("  ///////////////////////");
        System.out.println(" //// funny demo :) ////");
        System.out.println("///////////////////////");
        final JavaHeap<String, Integer> heap = new JavaHeap<>();

        heap.insert("Bad", 2);
        heap.insert("Mister", 1);
        heap.insert("Frosty", 3);

        System.out.println(heap.toDebugString());
        System.out.println();
        heap.insert("Bad", 3);

        System.out.println(heap.toDebugString());
        System.out.println();
    }

    private static void warmup() {
        System.out.println("  ////////////////////////////");
        System.out.println(" //// Warming the Frosty ////");
        System.out.println("////////////////////////////");

        System.out.println("[STATUS] Warming up...");

        final JavaHeap<Integer, Integer> javaHeap = new JavaHeap<>();
        final PriorityQueue<Integer> priorityQueue = new PriorityQueue<>();

        // Warmup insertion:
        for (int i = 0; i < WARMUP_SIZE / 2; ++i) {
            javaHeap.insert(i, i);
            priorityQueue.add(i);
        }

        // Warmup the update of JavaHeap:
        for (int i = 0; i < WARMUP_SIZE / 2; ++i) {
            javaHeap.insert(i, i);
            priorityQueue.add(i);
        }

        while (!priorityQueue.isEmpty()) {
            priorityQueue.remove();
        }

        while (!javaHeap.isEmpty()) {
            javaHeap.extractMinimum();
        }

        System.out.println("[STATUS] Warming up done!");
        System.out.println();
    }

    private static void benchmark() {
        System.out.println("  //////////////////////////");
        System.out.println(" //// Benchmarking :-) ////");
        System.out.println("//////////////////////////");

        final Integer[] dataArray = new Integer[BENCHMARK_SIZE];

        for (int i = 0; i < dataArray.length; ++i) {
            dataArray[i] = i % UNIQUE_VALUES;
        }

        shuffle(dataArray);

        final JavaHeap<Integer, Integer> javaHeap = new JavaHeap<>();
        final PriorityQueue<Integer> priorityQueue = new PriorityQueue<>();

        long startTime = System.currentTimeMillis();

        for (final Integer i : dataArray) {
            priorityQueue.add(i);
        }

        long endTime = System.currentTimeMillis();

        System.out.println("PriorityQueue.add() in " + (endTime - startTime) + 
                           " milliseconds.");

        startTime = System.currentTimeMillis();

        for (final Integer i : dataArray) {
            javaHeap.insert(i, i);
        }

        endTime = System.currentTimeMillis();

        System.out.println("JavaHeap.insert() in " + (endTime - startTime) + 
                           " milliseconds.");


        startTime = System.currentTimeMillis();

        while (!priorityQueue.isEmpty()) {
            priorityQueue.remove();
        }

        endTime = System.currentTimeMillis();

        System.out.println("PriorityQueue.remove() in " + 
                           (endTime - startTime) + " milliseconds.");

        startTime = System.currentTimeMillis();

        while (!javaHeap.isEmpty()) {
            javaHeap.extractMinimum();
        }

        endTime = System.currentTimeMillis();

        System.out.println("JavaHeap.extractMinimum() in " + (endTime - startTime) + 
                           " milliseconds.");


    }

    private static void shuffle(final Integer[] data) {
        final Random random = new Random();

        for (int i = 0; i < data.length; ++i) {
            final int j = random.nextInt(data.length);

            final Integer tmp = data[i];
            data[i] = data[j];
            data[j] = tmp;
        }
    }
}

JavaHeapTest.java:

package net.coderodde.util;

import java.util.NoSuchElementException;
import org.junit.Test;
import static org.junit.Assert.*;
import org.junit.Before;

public class JavaHeapTest {

    private JavaHeap<Integer, Float> heap = new JavaHeap<>();

    @Before
    public void before() {
        heap.clear();
    }

    @Test
    public void testInsert() {
        heap.insert(2, 2.0f);
        heap.insert(1, 1.0f);
        heap.insert(3, 3.0f);

        heap.insert(23, 3.0f);
        heap.insert(21, 1.0f);

        assertEquals((Integer) 1,  heap.extractMinimum());
        assertEquals((Integer) 21, heap.extractMinimum());
        assertEquals((Integer) 2,  heap.extractMinimum());
        assertEquals((Integer) 3,  heap.extractMinimum());
        assertEquals((Integer) 23, heap.extractMinimum());

        assertEquals(0, heap.size());

        heap.insert(1, 1.0f);

        assertEquals(1, heap.size());

        heap.insert(2, 2.0f);

        assertEquals(2, heap.size());

        heap.insert(1, 3.0f);

        assertEquals(3, heap.size());

        assertEquals((Integer) 1, heap.extractMinimum());
        assertEquals((Integer) 2, heap.extractMinimum());
        assertEquals((Integer) 1, heap.extractMinimum());
    }

    @Test(expected = NoSuchElementException.class)
    public void testExtractMinimumThrowsOnEmptyHeap() {
        heap.extractMinimum();
    }

    @Test
    public void testSize() {
        for (int i = 0; i < 10; ++i) {
            assertEquals(i, heap.size());
            heap.insert(i, 1.0f * i);
            assertEquals(i + 1, heap.size());
        }
    }

    @Test
    public void testIsEmpty() {
        assertTrue(heap.isEmpty());

        for (int i = 0; i < 10; ++i) {
            heap.insert(i, 1.0f * i);
            assertFalse(heap.isEmpty());
        }

        for (int i = 0; i < 10; ++i) {
            assertFalse(heap.isEmpty());
            heap.extractMinimum();
        }

        assertTrue(heap.isEmpty());
    }

    @Test
    public void testClear() {
        for (int i = 0; i < 100; ++i) {
            heap.insert(i, 2.0f * i);
            assertFalse(heap.isEmpty());
        }

        heap.clear();

        assertTrue(heap.isEmpty());

        heap.clear();

        assertTrue(heap.isEmpty());
    }
}

On sparse priority range (small number of distinct priorities compared to the size of the heap), my performance figures are as follows:


PriorityQueue.add() in 996 milliseconds.
JavaHeap.insert() in 2197 milliseconds.
PriorityQueue.remove() in 6966 milliseconds.
JavaHeap.extractMinimum() in 149 milliseconds.

As always, any critique is much appreciated.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.