Priority stack in Java

Question

I have this (simple) data structure that maintains multiple LIFO stacks. Each such stack is assigned a priority. The less is the value of a priority key, the higher the priority. Pushing operation requires the element to push and its priority key. Whenever popping the queue, it removes the most recently added element of the highest priority stack.

I have no clue where one would need such a data structures (maybe someone could tell me?), yet it was fun implementing it.

PriorityStack.java:

package net.coderodde.util;

import java.util.ArrayList;
import java.util.EmptyStackException;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

/**
 * This class implements a priority stack. Each element added comes with a 
 * priority key. The smaller the key, the higher is the priority of the element. 
 * The elements with the same priority are organized with LIFO rule 
 * ("last in, first out").
 * 
 * @author  Rodion "rodde" Efremov
 * @version 1.6
 * @param   <E> the element type.
 * @param   <P> the priority key type.
 */
public class PriorityStack<E, P extends Comparable<? super P>> {

    private final Map<P, List<E>> map;
    private int size;

    public PriorityStack() {
        this.map = new TreeMap<>();
    }

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

    /**
     * Pushes the input element to this stack.
     * 
     * @param element  the element to push.
     * @param priority the priority of the input element.
     */
    public void push(E element, P priority) {
        if (!map.containsKey(priority)) {
            map.put(priority, new ArrayList<>());
        }

        map.get(priority).add(element);
        ++size;
    }

    /**
     * Returns but not removes the topmost element of the highest priority 
     * substack.
     * 
     * @return the top element.
     */
    public E peek() {
        if (size == 0) {
            throw new EmptyStackException();
        }

        List<E> list = map.entrySet().iterator().next().getValue();
        return list.get(list.size() - 1);
    }

    /**
     * Pops the topmost element from the stack with the highest priority.
     * 
     * @return an element.
     * @throws EmptyStackException if the stack is empty.
     */
    public E pop() {
        if (size == 0) {
            throw new EmptyStackException();
        }

        Map.Entry<P, List<E>> entry = map.entrySet()
                                         .iterator()
                                         .next();
        P priority = entry.getKey();
        List<E> list = entry.getValue();
        E ret = list.remove(list.size() - 1);

        if (list.isEmpty()) {
            map.remove(priority);
        }

        --size;
        return ret;
    }

    public static void main(final String... args) {
        PriorityStack<String, Integer> stack = new PriorityStack<>();

        // -1: He is a 
        // 2:  funky programmer 
        // 3:  that likes to 
        // 5:  code algorithms 
        // 7:  and 
        // 8:  data structures! :-]

        stack.push("a", -1);
        stack.push("and", 7);
        stack.push("is", -1);
        stack.push("programmer", 2);
        stack.push("algorithms", 5);
        stack.push("code", 5);
        stack.push(":-]", 8);
        stack.push("He", -1);
        stack.push("to", 3);
        stack.push("structures!", 8);
        stack.push("likes", 3);
        stack.push("data", 8);
        stack.push("who", 3);
        stack.push("funky", 2);

        while (!stack.empty()) {
            System.out.print(stack.pop());
            System.out.print(" ");
        }

        System.out.println();
    }
}

So, what do you think?

Answer 1

You can consider implementing the Deque interface so that you can potentially open up its usage as an implementation of the Collections classes.

Your field size seems to be underused here, as I can sort of see a size() method using it too... BTW, regardless whether you use the Deque interface, your method empty() should be written as isEmpty(), to follow Java methods' naming convention as 'is...' for boolean return values.

Back to talking about Collections classes, they throw NoSuchElementException when no element can be returned, e.g. when empty, so perhaps you can throw that too instead of EmptyStackException.

Since you are relying on your TreeMap inner implementation, you can consider having your field declaration as TreeMap too instead of Map, this being one of those cases where using the implementation (class) rather than the interface makes more sense. With a TreeMap, you can replace map.entrySet().iterator().next().getValue() with the neater method map.firstEntry().

In your push() method, if you are on Java 8, you can consider computeIfAbsent():

public void push(E element, P priority) {
    map.computeIfAbsent(priority, ArrayList::new).add(element);
    size++;
}

Last but not least, you should have better unit testing too... Your current case looks OK already, so it's just a matter of comparing the output when you pop() to something you can compare with, such as a List. Remember to test for edge cases, such as getting from an empty instance, or adding null elements (which should be allowed given your ArrayList implementation).