Fixed heap removal bug.

[jagi.git] / src / jagi / event / Heap.java

package jagi.event;

import java.util.*;

public class Heap<V, K> {
    private static final Object[] empty = {};
    private final Comparator<? super K> cmp;
    private final Map<V, Integer> index = new IdentityHashMap<>();
    private Object[] vbuf = empty, kbuf = empty;
    private int size;

    public Heap(Comparator<? super K> cmp) {
	this.cmp = cmp;
    }

    @SuppressWarnings("unchecked")
    private V val(int i) {return((V)vbuf[i]);}
    @SuppressWarnings("unchecked")
    private K key(int i) {return((K)kbuf[i]);}

    private void raise(V val, K key, int i) {
	while(i > 0) {
	    int p = (i - 1) >>> 1;
	    if(cmp.compare(key(p), key) <= 0)
		break;
	    vbuf[i] = vbuf[p];
	    kbuf[i] = kbuf[p];
	    index.put(val(i), i);
	    i = p;
	}
	vbuf[i] = val;
	kbuf[i] = key;
	index.put(val, i);
    }

    private void lower(V val, K key, int i) {
	while(true) {
	    int c1 = (i << 1) + 1, c2 = c1 + 1;
	    if(c1 >= size)
		break;
	    int c = ((c2 < size) && (cmp.compare(key(c1), key(c2)) > 0)) ? c2 : c1;
	    if(cmp.compare(key(c), key) > 0)
		break;
	    vbuf[i] = vbuf[c];
	    kbuf[i] = kbuf[c];
	    index.put(val(i), i);
	    i = c;
	}
	vbuf[i] = val;
	kbuf[i] = key;
	index.put(val, i);
    }

    private void adjust(V val, K key, int i) {
	if((i > 0) && cmp.compare(key((i - 1) >> 1), key) > 0)
	    raise(val, key, i);
	else
	    lower(val, key, i);
    }

    public int size() {
	return(size);
    }

    public V peek() {
	return((size > 0) ? val(0) : null);
    }

    public V poll() {
	if(size == 0)
	    return(null);
	V ret = val(0);
	remove(0);
	index.remove(ret);
	return(ret);
    }

    public V remove() {
	if(size == 0)
	    throw(new NoSuchElementException());
	V ret = val(0);
	remove(0);
	index.remove(ret);
	return(ret);
    }

    public K keypeek() {
	return((size > 0) ? key(0) : null);
    }

    public void add(V val, K key) {
	if(index.containsKey(val))
	    throw(new IllegalStateException());
	int p = size++;
	if(p >= vbuf.length) {
	    int n = Math.max(vbuf.length * 2, 16);
	    vbuf = Arrays.copyOf(vbuf, n);
	    kbuf = Arrays.copyOf(kbuf, n);
	}
	raise(val, key, p);
    }

    public K update(V val, K key) {
	Integer p = index.get(val);
	if(p == null)
	    throw(new NoSuchElementException());
	K ret = key(p);
	adjust(val, key, p);
	return(ret);
    }

    public K set(V val, K key) {
	Integer p = index.get(val);
	if(p == null) {
	    add(val, key);
	    return(null);
	}
	K ret = key(p);
	adjust(val, key, p);
	return(ret);
    }

    private K remove(int p) {
	K ret = key(p);
	size--;
	if(p == size) {
	} else if(p < size) {
	    adjust(val(size), key(size), p);
	} else {
	    throw(new AssertionError());
	}
	vbuf[size] = null;
	kbuf[size] = null;
	return(ret);
    }

    public K remove(V val) {
	Integer p = index.remove(val);
	if(p == null)
	    return(null);
	return(remove(p));
    }

    public boolean contains(V val) {
	return(index.containsKey(val));
    }

    public String toString() {
	StringBuilder buf = new StringBuilder();
	buf.append('[');
	for(int i = 0; i < size; i++) {
	    if(i > 0)
		buf.append(", ");
	    buf.append(String.valueOf(kbuf[i]));
	    buf.append('=');
	    buf.append(String.valueOf(vbuf[i]));
	}
	buf.append(']');
	return(buf.toString());
    }
}