Probability data structure - algorithm

Question

The Idea is to have a data structure that you can access its elements only randomly, but based on a probability factor defined by the user for each element. So if the probability of the structure that contains 100 elements to yield x is 0.5, then, theoretically if we try to retrieve a random element a hundred times then x will be returned about \~50 times.
I couldn't find a ready-solution that does this, so this is my take on it:
import kotlin.math.absoluteValue
/**
*#author mhashim6 on 13/10/2019
*/
class ProbabilitySet<T>(private val items: Array<out Pair<T, Float>>) {
private var probabilityIndices: List<Int>
private fun calcFutureSize(count: Int, probability: Float) =
((count / (1f - probability)) - count).toInt().absoluteValue
init {
probabilityIndices = items.withIndex().flatMap { (i, item) ->
item.act { (_, probability) ->
calcFutureSize(items.size, probability).minus(items.size).act { delta ->
Iterable { ConstIterator(delta, i) }
}
}
}
}
fun next(): T = items.random().first
}
class ConstIterator(private var size: Int, private val const: Int) : IntIterator() {
override fun nextInt(): Int {
size--
return const
}
override fun hasNext(): Boolean = size > 0
}
fun <E> probabilitySetOf(vararg items: Pair<E, Float>) = ProbabilitySet(items)
inline fun <T, R> T.act(action: (T) -> R) = action(this)
I tried to make it mutable, but I met a lot of complexities regarding time and memory. So it's immutable for now.
Is this a viable Implementation?
Is there an implementation for this problem already?
How to make it mutable?

Answer 1

I assume that if the sum of elements' probabilities is not equal to 1, the actual element probability must be calculated by dividing its original probability by a sum of all elements' probabilities. For example, ProbabilitySet consisting of "A" to 0.1F and "B" to 0.3F returns "A" in 25% of cases and "B" in 75% of cases.
Here is my implementation of mutable ProbabilitySet with add running in O(1) and next running in O(logN):
class ProbabilitySet<E>(
private val random: Random = Random.Default
) {
private val nodes = mutableListOf<Node>()
private var sum = 0F
fun add(element: E, probability: Float) {
require(probability >= 0) { "[$element]'s probability ($probability) is less than 0" }
val oldSum = sum
sum += probability
nodes += Node(oldSum..sum, element)
}
fun isEmpty() = sum == 0F
fun next(): E {
if (isEmpty()) throw NoSuchElementException("ProbabilitySet is empty")
val index = random.nextFloat() * sum
return nodes[nodes.binarySearch {
when {
it.range.start > index -> 1
it.range.endInclusive < index -> -1
else -> 0
}
}].element
}
private inner class Node(
val range: ClosedRange<Float>,
val element: E
)
}
Factory method:
fun <E> probabilitySetOf(vararg items: Pair<E, Float>, random: Random = Random.Default) =
ProbabilitySet<E>(random).apply {
items.forEach { (element, probability) -> add(element, probability) }
}
Use case:
val set = probabilitySetOf("A" to 0.4F, "B" to 0.3F)
println(set.next())
set.add("C", 0.9F)
println(set.next())