aboutsummaryrefslogtreecommitdiff
path: root/src/main/java/gtPlusPlus/api/objects/data/WeightedCollection.java
blob: 5d990971695f48ab7d672072b6ab62c890f26a36 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
package gtPlusPlus.api.objects.data;

import java.util.Collection;
import java.util.Map;
import java.util.NavigableMap;
import java.util.Random;
import java.util.Set;
import java.util.TreeMap;

import gregtech.api.objects.XSTR;

public class WeightedCollection<E> implements Map<Integer, E> {

    private NavigableMap<Integer, E> map = new TreeMap<>();
    private Random random;
    private int total = 0;

    public WeightedCollection() {
        this(new XSTR());
    }

    public WeightedCollection(Random random) {
        this.random = random;
    }

    public E add(int weight, E object) {
        if (weight <= 0) return null;
        total += weight;
        return map.put(total, object);
    }

    private E next() {
        int value = random.nextInt(total) + 1; // Can also use floating-point weights
        return map.ceilingEntry(value)
            .getValue();
    }

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

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

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

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

    public E get() {
        return next();
    }

    @Override
    public E get(Object key) {
        return next();
    }

    @Override
    public void putAll(Map m) {
        map.putAll(m);
    }

    @Override
    public void clear() {
        map.clear();
        this.total = 0;
    }

    @Override
    public Set keySet() {
        return map.keySet();
    }

    @Override
    public Collection values() {
        return map.values();
    }

    @Override
    public Set entrySet() {
        return map.entrySet();
    }

    @Override
    public E put(Integer key, E value) {
        return add(key, value);
    }

    @Override
    public E remove(Object key) {
        return map.remove(key);
    }
}