[Java] 要素追加時の優先度設定と優先度の更新が可能な優先度つきキュー

SettablePriorityQueue.java

import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Queue;

/**
 * 優先度の設定が可能なキューです。<br>
 * JDKの{@link PriorityQueue}とは異なり要素を優先度つきでキューに追加することができます。<br>
 * 後から優先度を更新することも可能です。<br><br>
 * 
 * 優先度つき要素の追加は{@link #offer(E, Priority)}か {@link #add(E, Priority)}で行えます。<br>
 * また優先度の更新は{@link #updatePriority(E, Priority)}で行えます。<br><br>
 * 
 * {@link #offer(E)},{@link #add(E)}で要素を追加した場合、要素の順序付けは、要素自身の自然順序付けに従って行われます。<br>
 * これは{@link PriorityQueue}と同等の動きとなります。
 * そのため優先度つきで要素を追加しない場合は、このクラスではなく{@link PriorityQueue}の使用を推奨します。<br>
 * {@link #offer(E, Priority)}と{@link #offer(E)}のどちらも使った場合の順序付けも、
 * 要素自身の自然順序付けに従って行われます。<br>
 * ({@link #offer(E, Priority)}で指定した優先度は無視されます。)<br>
 * すなわち要素すべてに対して優先度が設定されない限り常に要素自身の自然順序付けが使用されます。<br>
 * <br>
 * サンプル
 * <pre>
 * SettablePriorityQueue&lt;String, Integer&gt; solveOrderQueue 
 *     = new SettablePriorityQueue&lt;String, Integer&gt;();
 *     
 * solveOrderQueue.offer("簡単な問題", 5);
 * solveOrderQueue.offer("普通の問題", 6);
 * solveOrderQueue.offer("難しい問題", 7);
 *       
 * System.out.println(solveOrderQueue); // -> [簡単な問題, 普通の問題, 難しい問題]
 *
 * // 優先度の更新
 * solveOrderQueue.updatePriority("難しい問題", 1);
 *       
 * System.out.println(solveOrderQueue); // -> [難しい問題, 普通の問題, 簡単な問題]
 * </pre>
 * 
 * @param <E> コレクション内に存在する要素の型
 * @param <Priority> 優先度の型 {@link Comparable}を実装していない型を指定する場合、
 *                    コンストラクタに Comparator<Priority> を渡してください。
 * @author kenji suzuki
 */
public class SettablePriorityQueue<E, Priority> implements Queue<E> {

    protected static final int DEFAULT_INITIAL_CAPACITY = 11;

    /** 処理を委譲するキュー */
    protected PriorityQueue<E> queue;

    /** キューの要素に対応した優先度を持つマップ */
    protected Map<E, Priority> priorityMap;
    
    /** 優先度の{@link Comparator} */
    protected Comparator<Priority> priorityComparator = new DefaultPriorityComparator();

    /**
     * デフォルトの初期容量 (11) を持つ {@link SettablePriorityQueue} を作成します。
     */
    public SettablePriorityQueue() {
        this(DEFAULT_INITIAL_CAPACITY);
    }

    /**
     * 指定された初期容量を持つ {@link SettablePriorityQueue} を作成します。
     * @param initialCapacity この優先度キューの初期容量 
     */
    public SettablePriorityQueue(int initialCapacity) {
        this(initialCapacity, null);
    }
    
    /**
     * 指定されたコンパレータに従って要素を順序付けする{@link SettablePriorityQueue} を作成します。
     * @param priorityComparator この優先度キューを順序付けするために使用されるコンパレータ。
     *                            null の場合、型パラメータ&lt;Priority&gt;の自然順序付けが使用される。
     */
    public SettablePriorityQueue(Comparator<Priority> priorityComparator) {
        this(DEFAULT_INITIAL_CAPACITY, priorityComparator);
    }
    
    /**
     * 指定されたコンパレータに従って要素を順序付けする、
     * 指定された初期容量を持つ {@link SettablePriorityQueue} を作成します。 
     * @param initialCapacity この優先度キューの初期容量
     * @param priorityComparator この優先度キューを順序付けするために使用されるコンパレータ。
     *                            null の場合、型パラメータ&lt;Priority&gt;の自然順序付けが使用される。
     */
    public SettablePriorityQueue(int initialCapacity, Comparator<Priority> priorityComparator) {
        queue = new PriorityQueue<>(initialCapacity, new QueueComparator());
        priorityMap = new HashMap<>(initialCapacity);
        if (priorityComparator != null) {
            this.priorityComparator = priorityComparator;
        }
    }

    /**
     * 指定された要素を、指定された優先度で、このキューに挿入します。
     * @param e 追加する要素
     * @param priority 追加する要素の優先度
     * @return true ({@link Queue#offer(E)}で指定されているとおり) 
     * @exception NullPointerException 指定された優先度が null である場合
     */
    public boolean offer(E e, Priority priority) {
        checkNotNull(priority);
        priorityMap.put(e, priority);
        return queue.offer(e);
    }

    /**
     * 指定された要素を、指定された優先度で、このキューに挿入します。
     * @param e 追加する要素
     * @param priority 追加する要素の優先度
     * @return true ({@link Collection#add(E)}で指定されているとおり) 
     * @exception NullPointerException 指定された優先度が null である場合
     */
    public void add(E e, Priority priority) {
        offer(e, priority);
    }

    /**
     * 指定された要素の優先度を、指定された優先度で更新します。
     * @param e 優先度を更新する要素
     * @param priority 要素の優先度
     */
    public void updatePriority(E e, Priority priority) {
        if (priorityMap.containsKey(e)) {
            remove(e);
            offer(e, priority);
        } else {
            // 優先度が未設定の要素に対し、優先度が設定された場合再構築が必要。
            // (offerを呼び出すだけでは、正しい順序になる保証がない)
            priorityMap.put(e, priority);
            rebuild();
        }
    }

    /**
     * 指定された要素の優先度を返します。
     * @param e 関連付けられた優先度が返される要素
     * @return 要素に関連付けられた優先度
     */
    public Priority getPriority(E e) {
        return priorityMap.get(e);
    }

    /**
     * 指定された要素をこの優先度キューに挿入します。
     * {@link #add(E, Priority)}ではなく、こちらを使って要素を挿入した場合、
     * 要素の順序付けには要素自身の自然順序付けが使用されます。
     * すなわち{@link PriorityQueue}と同じ順序付けになります。
     * そのためこの場合、追加する要素は{@link Comparable}を実装している必要があります。
     * @param e 追加する要素
     * @return true ({@link Collection#add(E)}で指定されているとおり) 
     * @exception NullPointerException 指定された要素が null である場合
     * @exception ClassCastException 指定された要素が{@link Comparable}を実装していない場合
     *                                 (2回目以降のoffer(E)で発生します。) 
     */
    @Override
    public boolean add(E e) {
        return offer(e);
    }
    
    /**
     * 指定された要素をこの優先度キューに挿入します。
     * {@link #add(E, Priority)}ではなく、こちらを使って要素を挿入した場合、
     * 要素の順序付けには要素自身の自然順序付けが使用されます。
     * すなわち{@link PriorityQueue}と同じ順序付けになります。
     * そのためこの場合、追加する要素は{@link Comparable}を実装している必要があります。
     * @param e 追加する要素
     * @return true ({@link Queue#offer(E)}で指定されているとおり) 
     * @exception NullPointerException 指定された要素が null である場合
     * @exception ClassCastException 指定された要素が{@link Comparable}を実装していない場合
     *                                 (2回目以降のoffer(E)で発生します。)
     */
    @Override
    public boolean offer(E e) {
        checkNotNull(e);
        return queue.offer(e);
    }

    @Override
    public int size() {
        return queue.size();
    }

    @Override
    public boolean isEmpty() {
        return queue.isEmpty();
    }

    @Override
    public boolean contains(Object o) {
        return queue.contains(o);
    }

    @Override
    public E remove() {
        return poll();
    }

    @Override
    public E poll() {
        E e = queue.poll();
        priorityMap.remove(e);
        return e;
    }

    @Override
    public Object[] toArray() {
        return queue.toArray();
    }

    @Override
    public E element() {
        return queue.element();
    }

    @Override
    public E peek() {
        return queue.peek();
    }
    
    /**
     * このキュー内の要素の反復子を返します。
     * {@link PriorityQueue#iterator()}と同様に反復子が要素を返す特定の順序はありません。
     * @return キュー内の要素の反復子
     */
    @Override
    public Iterator<E> iterator() {
        return queue.iterator();
    }

    @Override
    public <T> T[] toArray(T[] a) {
        return queue.toArray(a);
    }

    @Override
    public boolean remove(Object o) {
        priorityMap.remove(o);
        return queue.remove(o);
    }

    @Override
    public boolean containsAll(Collection<?> c) {
        return queue.containsAll(c);
    }

    @Override
    public boolean addAll(Collection<? extends E> c) {
        return queue.addAll(c);
    }

    @Override
    public boolean removeAll(Collection<?> c) {
        boolean removeAll = queue.removeAll(c);
        if (removeAll) {
            for (Object object : c) {
                priorityMap.remove(object);
            }
            rebuild();
        }
        return removeAll;
    }

    @Override
    public boolean retainAll(Collection<?> c) {
        return queue.retainAll(c);
    }

    @Override
    public void clear() {
        queue.clear();
        priorityMap.clear();
    }

    @Override
    public boolean equals(Object o) {
        return queue.equals(o);
    }

    @Override
    public int hashCode() {
        return queue.hashCode();
    }

    /**
     * このキューの文字列表現を返します。
     * 文字列表現は、キューの要素を順序付けに従って取り出した順に角括弧 ("[]") で囲んで示すリストです。
     * 隣接する要素は、文字 ", " (カンマと空白文字) によって区切られます。
     * 各要素は、String.valueOf(Object) を実行したかのように文字列に変換されます。
     * @return このキューの文字列表現
     */
    @Override
    public String toString() {
        // PriorityQueue#toStringは順序を保証しないので単純に委譲はできない
        PriorityQueue<E> q = new PriorityQueue<>(queue);
        StringBuilder sb = new StringBuilder("[");
        while (!q.isEmpty()) {
            sb.append(q.poll());
            if (0 < q.size()) {
                sb.append(", ");
            }
        }
        return sb.append("]").toString();
    }
    
    protected void checkNotNull(Object object) {
        if (object == null) {
            throw new NullPointerException();
        }
    }
    
    protected void rebuild() {
        Object[] array = queue.toArray();
        queue.clear();
        for (Object element : array) {
            queue.offer((E) element);
        }
    }
        
    protected class QueueComparator implements Comparator<E> {

        @Override
        public int compare(E e1, E e2) {
            Priority p1 = priorityMap.get(e1);
            Priority p2 = priorityMap.get(e2);
            // どちらかの優先度の設定がない場合、要素同士を比較
            if (p1 == null || p2 == null) {
                return ((Comparable<E>) e1).compareTo(e2);
            }
            return priorityComparator.compare(p1, p2);
        }

    }
    
    protected class DefaultPriorityComparator implements Comparator<Priority> {

        @Override
        public int compare(Priority p1, Priority p2) {
            return ((Comparable<Priority>) p1).compareTo(p2);
        }

    }

}