[FIXED] Speed up random weighted choice without replacement in python

Issue

I want to sample ~10⁷ times from a population of ~10⁷ integers without replacements and with weights, each time picking 10 elements. After each sampling I change the weights. I have timed two approaches (python3 and numpy) in the following script. Both approaches seem painfully slow to me, do you see a way of speeding it up?

import numpy as np
import random

@profile
def test_choices():
    population = list(range(10**7))
    weights = np.random.uniform(size=10**7)
    np_weights = np.array(weights)

    def numpy_choice():
        np_w = np_weights / sum(np_weights)
        c = np.random.choice(population, size=10, replace=False, p=np_w)

    def python_choice():
        c = []
        while len(c) < 10:
            c += random.choices(population=population, weights=weights, k=10 - len(c))
            c = list(set(c))

    for i in range(10**1):

        numpy_choice()
        python_choice()

        add_weight = np.random.uniform()
        random_element = random.randint(0, 10**7)
        weights[random_element] += add_weight
        np_weights[random_element] += add_weight


test_choices()

With the timer result:

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    24        10   20720062.0 2072006.2     56.6          numpy_choice()
    25        10   15593925.0 1559392.5     42.6          python_choice()

Solution

This is just a comment on jdhesa's answer. The question was if it is useful to consider the case where only one weight is incresed -> Yes it is!

Example

@nb.njit(parallel=True)
def numba_choice_opt(population, weights, k,wc,b_full_wc_calc,ind,value):
    # Get cumulative weights
    if b_full_wc_calc:
        acc=0
        for i in range(weights.shape[0]):
            acc+=weights[i]
            wc[i]=acc
    #Increase only one weight (faster than recalculating the cumulative  weight)
    else:
        weights[ind]+=value
        for i in nb.prange(ind,wc.shape[0]):
            wc[i]+=value

    # Total of weights
    m = wc[-1]
    # Arrays of sample and sampled indices
    sample = np.empty(k, population.dtype)
    sample_idx = np.full(k, -1, np.int32)
    # Sampling loop
    i = 0
    while i < k:
        # Pick random weight value
        r = m * np.random.rand()
        # Get corresponding index
        idx = np.searchsorted(wc, r, side='right')
        # Check index was not selected before
        # If not using Numba you can just do `np.isin(idx, sample_idx)`
        for j in range(i):
            if sample_idx[j] == idx:
                continue
        # Save sampled value and index
        sample[i] = population[idx]
        sample_idx[i] = population[idx]
        i += 1
    return sample

Example

np.random.seed(0)
population = np.random.randint(100, size=1_000_000)
weights = np.random.rand(len(population))
k = 10
wc = np.empty_like(weights)

#Initial calculation 
%timeit numba_choice_opt(population, weights, k,wc,True,0,0)
#1.41 ms ± 9.21 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

#Increase weight[100] by 3 and calculate
%timeit numba_choice_opt(population, weights, k,wc,False,100,3)
#213 µs ± 6.06 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

#For comparison
#Please note that it is the memory allcocation of wc which makes
#it so much slower than the initial calculation from above
%timeit numba_choice(population, weights, k)
#4.23 ms ± 64.9 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Answered By - max9111