Fixed heap removal bug.
[jagi.git] / src / jagi / event / Heap.java
1 package jagi.event;
2
3 import java.util.*;
4
5 public class Heap<V, K> {
6     private static final Object[] empty = {};
7     private final Comparator<? super K> cmp;
8     private final Map<V, Integer> index = new IdentityHashMap<>();
9     private Object[] vbuf = empty, kbuf = empty;
10     private int size;
11
12     public Heap(Comparator<? super K> cmp) {
13         this.cmp = cmp;
14     }
15
16     @SuppressWarnings("unchecked")
17     private V val(int i) {return((V)vbuf[i]);}
18     @SuppressWarnings("unchecked")
19     private K key(int i) {return((K)kbuf[i]);}
20
21     private void raise(V val, K key, int i) {
22         while(i > 0) {
23             int p = (i - 1) >>> 1;
24             if(cmp.compare(key(p), key) <= 0)
25                 break;
26             vbuf[i] = vbuf[p];
27             kbuf[i] = kbuf[p];
28             index.put(val(i), i);
29             i = p;
30         }
31         vbuf[i] = val;
32         kbuf[i] = key;
33         index.put(val, i);
34     }
35
36     private void lower(V val, K key, int i) {
37         while(true) {
38             int c1 = (i << 1) + 1, c2 = c1 + 1;
39             if(c1 >= size)
40                 break;
41             int c = ((c2 < size) && (cmp.compare(key(c1), key(c2)) > 0)) ? c2 : c1;
42             if(cmp.compare(key(c), key) > 0)
43                 break;
44             vbuf[i] = vbuf[c];
45             kbuf[i] = kbuf[c];
46             index.put(val(i), i);
47             i = c;
48         }
49         vbuf[i] = val;
50         kbuf[i] = key;
51         index.put(val, i);
52     }
53
54     private void adjust(V val, K key, int i) {
55         if((i > 0) && cmp.compare(key((i - 1) >> 1), key) > 0)
56             raise(val, key, i);
57         else
58             lower(val, key, i);
59     }
60
61     public int size() {
62         return(size);
63     }
64
65     public V peek() {
66         return((size > 0) ? val(0) : null);
67     }
68
69     public V poll() {
70         if(size == 0)
71             return(null);
72         V ret = val(0);
73         remove(0);
74         index.remove(ret);
75         return(ret);
76     }
77
78     public V remove() {
79         if(size == 0)
80             throw(new NoSuchElementException());
81         V ret = val(0);
82         remove(0);
83         index.remove(ret);
84         return(ret);
85     }
86
87     public K keypeek() {
88         return((size > 0) ? key(0) : null);
89     }
90
91     public void add(V val, K key) {
92         if(index.containsKey(val))
93             throw(new IllegalStateException());
94         int p = size++;
95         if(p >= vbuf.length) {
96             int n = Math.max(vbuf.length * 2, 16);
97             vbuf = Arrays.copyOf(vbuf, n);
98             kbuf = Arrays.copyOf(kbuf, n);
99         }
100         raise(val, key, p);
101     }
102
103     public K update(V val, K key) {
104         Integer p = index.get(val);
105         if(p == null)
106             throw(new NoSuchElementException());
107         K ret = key(p);
108         adjust(val, key, p);
109         return(ret);
110     }
111
112     public K set(V val, K key) {
113         Integer p = index.get(val);
114         if(p == null) {
115             add(val, key);
116             return(null);
117         }
118         K ret = key(p);
119         adjust(val, key, p);
120         return(ret);
121     }
122
123     private K remove(int p) {
124         K ret = key(p);
125         size--;
126         if(p == size) {
127         } else if(p < size) {
128             adjust(val(size), key(size), p);
129         } else {
130             throw(new AssertionError());
131         }
132         vbuf[size] = null;
133         kbuf[size] = null;
134         return(ret);
135     }
136
137     public K remove(V val) {
138         Integer p = index.remove(val);
139         if(p == null)
140             return(null);
141         return(remove(p));
142     }
143
144     public boolean contains(V val) {
145         return(index.containsKey(val));
146     }
147
148     public String toString() {
149         StringBuilder buf = new StringBuilder();
150         buf.append('[');
151         for(int i = 0; i < size; i++) {
152             if(i > 0)
153                 buf.append(", ");
154             buf.append(String.valueOf(kbuf[i]));
155             buf.append('=');
156             buf.append(String.valueOf(vbuf[i]));
157         }
158         buf.append(']');
159         return(buf.toString());
160     }
161 }