WeightedSet.java

package net.wouto.simplemongo;

import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Random;
import java.util.Set;

public class WeightedSet<T> {
    
    private HashMap<T, Integer> map = new HashMap();
    private static Random RANDOM = new Random();
    private int weight;

    public T getRandom() {
        int sel = RANDOM.nextInt(weight);
        T obj = null;
        for (Entry<T, Integer> e : map.entrySet()) {
            sel -= e.getValue();
            if (sel <= 0) {
                obj = e.getKey();
                break;
            }
        }
        return obj;
    }
    
    public int getSetWeight() {
        return this.weight;
    }
    
    public int size() {
        return map.size();
    }

    public boolean isEmpty() {
        return map.isEmpty();
    }

    public boolean contains(Object key) {
        return map.containsKey(key);
    }

    public boolean containsWeight(Object value) {
        return map.containsValue(value);
    }

    public Integer getWeight(Object key) {
        return map.get(key);
    }
    
    public Integer add(T key, Integer value) {
        this.map.put(key, value);
        this.weight += value;
        return value;
    }

    public Integer remove(Object key) {
        Integer i = map.remove(key);
        this.weight -= i;
        return i;
    }

    public void putAll(Map<? extends T, ? extends Integer> m) {
        map.putAll(m);
        for (Integer j : m.values()) {
            this.weight += j;
        }
    }

    public void clear() {
        map.clear();
        this.weight = 0;
    }

    public Set<T> getUnweightedSet() {
        return map.keySet();
    }

    public Collection<Integer> getWeightSet() {
        return map.values();
    }

    public Set<Map.Entry<T, Integer>> entrySet() {
        return map.entrySet();
    }
    
}