lewysDavies / Java-Probability-Collection

* Copyright (c) 2020 Lewys Davies

* 

* Permission is hereby granted, free of charge, to any person obtaining a copy

* of this software and associated documentation files (the "Software"), to deal

* in the Software without restriction, including without limitation the rights

* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

* copies of the Software, and to permit persons to whom the Software is

* furnished to do so, subject to the following conditions:

*

* The above copyright notice and this permission notice shall be included in all

* copies or substantial portions of the Software.

* 

* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE

* SOFTWARE.

*/

package com.lewdev.probabilitylib;

import java.util.Comparator;

import java.util.Iterator;

import java.util.NavigableSet;

import java.util.Objects;

import java.util.SplittableRandom;

import java.util.TreeSet;

/**

 * ProbabilityCollection for retrieving random elements based on probability.

 * <br>

 * <br>

 * <b>Selection Algorithm Implementation</b>:

 * <p>

 * <ul>

 * <li>Elements have a "block" of space, sized based on their probability share

 * <li>"Blocks" start from index 1 and end at the total probability of all elements

 * <li>A random number is selected between 1 and the total probability

 * <li>Which "block" the random number falls in is the element that is selected

 * <li>Therefore "block"s with larger probability have a greater chance of being

 * selected than those with smaller probability.

 * </p>

 * </ul>

 * 

 * @author Lewys Davies

 * @version 0.7

 *

 * @param <E> Type of elements

 */

public class ProbabilityCollection<E> {

	private final NavigableSet<ProbabilitySetElement<E>> collection;

	private final SplittableRandom random = new SplittableRandom();

	private int totalProbability;

	/**

	 * Construct a new Probability Collection

	 */

	public ProbabilityCollection() {

		this.collection = new TreeSet<>(Comparator.comparingInt(ProbabilitySetElement::getIndex));

		this.totalProbability = 0;

	}

	/**

	 * @return Number of objects inside the collection

	 */

	public int size() {

		return this.collection.size();

	}

	/**

	 * @return True if collection contains no elements, else False

	 */

	public boolean isEmpty() {

		return this.collection.isEmpty();

	}

	/**

	 * @param <E> object

	 * @return True if collection contains the object, else False

	 * @throws IllegalArgumentException if object is null

	 */

	public boolean contains(E object) {

		if(object == null) {

			throw new IllegalArgumentException("Cannot check if null object is contained in this collection");

		}

		return this.collection.stream()

			.anyMatch(entry -> entry.getObject().equals(object));

	}

	/**

	 * @return Iterator over this collection

	 */

	public Iterator<ProbabilitySetElement<E>> iterator() {

		return this.collection.iterator();

	}

	/**

	 * Add an object to this collection

	 * 

	 * @param <E> object. Not null.

	 * @param probability share. Must be greater than 0.

	 * 

	 * @throws IllegalArgumentException if object is null

	 * @throws IllegalArgumentException if probability <= 0

	 */

	public void add(E object, int probability) {

		if(object == null) {

			throw new IllegalArgumentException("Cannot add null object");

		}

		if(probability <= 0) {

			throw new IllegalArgumentException("Probability must be greater than 0");

		}

		this.collection.add(new ProbabilitySetElement<E>(object, probability));

		this.totalProbability += probability;

		this.updateIndexes();

	}

	/**

	 * Remove a object from this collection

	 * 

	 * @param <E> object

	 * @return True if object was removed, else False.

	 * 

	 * @throws IllegalArgumentException if object is null

	 */

	public boolean remove(E object) {

		if(object == null) {

			throw new IllegalArgumentException("Cannot remove null object");

		}

		Iterator<ProbabilitySetElement<E>> it = this.iterator();

		boolean removed = it.hasNext();

		while(it.hasNext()) {

			ProbabilitySetElement<E> entry = it.next();

			if(entry.getObject().equals(object)) {

				this.totalProbability -= entry.getProbability();

				it.remove();

			}

		}

		this.updateIndexes();

		return removed;

	}

	/**

	 * Remove all objects from this collection

	 */

	/**

	 * Get a random object from this collection, based on probability.

	 * 

	 * @return <E> Random object

	 * 

	 * @throws IllegalStateException if this collection is empty

	 */

	public E get() {

		if(this.isEmpty()) {

			throw new IllegalStateException("Cannot get an object out of a empty collection");

		}

		ProbabilitySetElement<E> toFind = new ProbabilitySetElement<>(null, 0);

		toFind.setIndex(this.random.nextInt(1, this.totalProbability + 1));

		return Objects.requireNonNull(this.collection.floor(toFind).getObject());

	}

	/**

	 * @return Sum of all element's probability

	 */

	public final int getTotalProbability() {

		return this.totalProbability;

	}

	/*

	 * Calculate the size of all element's "block" of space: 

	 * i.e 1-5, 6-10, 11-14, 15, 16

	 * 

	 * We then only need to store the start index of each element,

	 * as we make use of the TreeSet#floor

	 */

	private final void updateIndexes() {

		int previousIndex = 0;

		for(ProbabilitySetElement<E> entry : this.collection) {

			previousIndex = entry.setIndex(previousIndex + 1) + (entry.getProbability() - 1);

		}

	}

	/**

	 * Used internally to store information about a object's

	 * state in a collection. Specifically, the probability 

	 * and index within the collection.

	 * 

	 * Indexes refer to the start position of this element's "block" of space.

	 * The space between element "block"s represents their probability of being selected

	 * 

	 * @author Lewys Davies

	 *

	 * @param <T> Type of element

	 */

	final static class ProbabilitySetElement<T> {

		private final T object;

		private final int probability;

		private int index;

		/**

		 * @param <T> object

		 * @param probability

		 */

		protected ProbabilitySetElement(T object, int probability) {

			this.object = object;

			this.probability = probability;

		}

		/**

		 * @return <T> The actual object

		 */

		public final T getObject() {

			return this.object;

		}

		/**

		 * @return Probability share in this collection

		 */

		public final int getProbability() {

			return this.probability;

		}

		// Used internally, see this class's documentation

		private final int getIndex() {

			return this.index;

		}

		// Used Internally, see this class's documentation

		private final int setIndex(int index) {

			this.index = index;

			return this.index;

		}

	}

}