133 lines
3.7 KiB
Java
133 lines
3.7 KiB
Java
package com.hypixel.hytale.builtin.hytalegenerator.datastructures;
|
|
|
|
import java.util.ArrayList;
|
|
import java.util.HashMap;
|
|
import java.util.HashSet;
|
|
import java.util.List;
|
|
import java.util.Map;
|
|
import java.util.Random;
|
|
import java.util.Set;
|
|
import java.util.function.BiConsumer;
|
|
import javax.annotation.Nonnull;
|
|
|
|
public class WeightedMap<T> {
|
|
@Nonnull
|
|
private final Set<T> elementSet;
|
|
@Nonnull
|
|
private final List<T> elements;
|
|
@Nonnull
|
|
private final List<Double> weights;
|
|
@Nonnull
|
|
private final Map<T, Integer> indices;
|
|
private double totalWeight = 0.0;
|
|
private boolean immutable = false;
|
|
|
|
public WeightedMap(@Nonnull WeightedMap<T> other) {
|
|
this.totalWeight = other.totalWeight;
|
|
this.elementSet = new HashSet<>(other.elementSet);
|
|
this.elements = new ArrayList<>(other.elements);
|
|
this.weights = new ArrayList<>(other.weights);
|
|
this.indices = new HashMap<>(other.indices);
|
|
this.immutable = other.immutable;
|
|
}
|
|
|
|
public WeightedMap() {
|
|
this(2);
|
|
}
|
|
|
|
public WeightedMap(int initialCapacity) {
|
|
this.elementSet = new HashSet<>(initialCapacity);
|
|
this.elements = new ArrayList<>(initialCapacity);
|
|
this.weights = new ArrayList<>(initialCapacity);
|
|
this.indices = new HashMap<>(initialCapacity);
|
|
}
|
|
|
|
@Nonnull
|
|
public WeightedMap<T> add(@Nonnull T element, double weight) {
|
|
if (element == null) {
|
|
throw new NullPointerException();
|
|
} else if (this.immutable) {
|
|
throw new IllegalStateException("method can't be called when object is immutable");
|
|
} else if (weight < 0.0) {
|
|
throw new IllegalArgumentException("weight must be positive");
|
|
} else {
|
|
this.elements.add(element);
|
|
this.weights.add(weight);
|
|
this.elementSet.add(element);
|
|
this.totalWeight += weight;
|
|
this.indices.put(element, this.indices.size());
|
|
return this;
|
|
}
|
|
}
|
|
|
|
public double get(@Nonnull T element) {
|
|
if (element == null) {
|
|
throw new NullPointerException();
|
|
} else if (this.immutable) {
|
|
throw new IllegalStateException("method can't be called when object is immutable");
|
|
} else {
|
|
return !this.elementSet.contains(element) ? 0.0 : this.weights.get(this.indices.get(element));
|
|
}
|
|
}
|
|
|
|
public T pick(@Nonnull Random rand) {
|
|
if (rand == null) {
|
|
throw new NullPointerException();
|
|
} else if (this.elements.isEmpty()) {
|
|
throw new IllegalStateException("can't be empty when calling this method");
|
|
} else {
|
|
double pointer = rand.nextDouble() * this.totalWeight;
|
|
|
|
for (int i = 0; i < this.elements.size(); i++) {
|
|
pointer -= this.weights.get(i);
|
|
if (pointer <= 0.0) {
|
|
return this.elements.get(i);
|
|
}
|
|
}
|
|
|
|
return this.elements.getLast();
|
|
}
|
|
}
|
|
|
|
public int size() {
|
|
return this.elements.size();
|
|
}
|
|
|
|
@Nonnull
|
|
public List<T> allElements() {
|
|
return new ArrayList<>(this.elements);
|
|
}
|
|
|
|
public void makeImmutable() {
|
|
this.immutable = true;
|
|
}
|
|
|
|
public boolean isImmutable() {
|
|
return this.immutable;
|
|
}
|
|
|
|
public void forEach(@Nonnull BiConsumer<T, Double> consumer) {
|
|
for (int i = 0; i < this.elements.size(); i++) {
|
|
consumer.accept(this.elements.get(i), this.weights.get(i));
|
|
}
|
|
}
|
|
|
|
@Nonnull
|
|
@Override
|
|
public String toString() {
|
|
return "WeighedMap{elementSet="
|
|
+ this.elementSet
|
|
+ ", elements="
|
|
+ this.elements
|
|
+ ", weights="
|
|
+ this.weights
|
|
+ ", indices="
|
|
+ this.indices
|
|
+ ", totalWeight="
|
|
+ this.totalWeight
|
|
+ ", immutable="
|
|
+ this.immutable
|
|
+ "}";
|
|
}
|
|
}
|