Skip to content

Instantly share code, notes, and snippets.

@arafsheikh
Created July 1, 2022 17:17
Show Gist options
  • Save arafsheikh/21270b137c80524f1f3d29651099f748 to your computer and use it in GitHub Desktop.
Save arafsheikh/21270b137c80524f1f3d29651099f748 to your computer and use it in GitHub Desktop.
TreeMultiMap: Data structure that (partially) implements a TreeMap with a key and a list of values
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
public class TreeMultiMap<K extends Comparable<K>, V> {
private final TreeMap<K, List<V>> treeMap;
int size;
public TreeMultiMap() {
treeMap = new TreeMap<>(Comparable::compareTo);
size = 0;
}
public void put(K key, V value) {
if (treeMap.containsKey(key)) {
treeMap.get(key).add(value);
} else {
treeMap.put(key, new ArrayList<>() {{
add(value);
}});
}
size++;
}
public V removeFirst() {
K key = treeMap.firstKey();
if (key == null) {
return null;
}
size--;
if (treeMap.get(key).size() == 1) {
V value = treeMap.remove(key).get(0);
treeMap.remove(key);
return value;
} else {
return treeMap.get(key).remove(0);
}
}
public List<Map.Entry<K, V>> getAllDescending() {
List<Map.Entry<K, V>> entries = new ArrayList<>();
for (Map.Entry<K, List<V>> entry : treeMap.descendingMap().entrySet()) {
for (V value : entry.getValue()) {
entries.add(new Map.Entry<>() {
@Override
public K getKey() {
return entry.getKey();
}
@Override
public V getValue() {
return value;
}
@Override
public V setValue(V value) {
throw new UnsupportedOperationException();
}
});
}
}
return entries;
}
public int size() {
return size;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment