One-pass sampling without replacement by Efraimidis and Spirakis algorithm

one_pass_sampling_without_replacement.py

# Efraimidis and Spirakis algorithm
# https://www.sciencedirect.com/science/article/abs/pii/S002001900500298
import numpy as np
from numba import njit


@njit(nogil=True)
def one_pass_sampling_without_replacement(n, k, weights):
    # Draw a uniform random variable for each item
    u = np.random.rand(n)

    # Compute a random variable X for each item
    Z = np.log(u) / weights

    # Find the indices of the k largest values
    indices = np.argpartition(Z, -k)[-k:]

    # Return the indices
    return indices