Iterate through a subset of a Cartesian product where all elements are selected (near-)equally

Viewed 219

I have a large set of samples which can be described by three parameters (let's call them a, b, and c), for instance in a tuple (a, b, c), which each can have a finite number of values. For example, there are 25 possible values for a (indexed 0..24), 20 possible values for b, and 3 possible values for c. Every combination of a, b, and c is represented in this dataset, so in this example, my dataset has 1500 samples (25 × 20 × 3).

I want to randomly select a subset of n samples from this dataset (without duplicates). However, this random sample must have the property that all possible values for a, b, and c are represented equally (or as close to equally as possible, if the number of selected samples is not divisible by the number of possible values for the parameters).

For example, if I select 100 samples, I want each value of a to be represented 4 times, each value of b to be represented 5 times, and each value of c to be represented 33 times (one value can be represented 34 times to satisfy the total number of samples selected, and it doesn't matter which value this is). I do not care about the exact combinations of (a, b, c), as long as the total number of occurrences of each parameter value is correct.

My current implementation is as follows:

import random

n_a = 25
n_b = 20
n_c = 3

n_desired = 100

# generate random ordering for selections
order_a = random.sample(range(n_a), k=n_a)
order_b = random.sample(range(n_b), k=n_b)
order_c = random.sample(range(n_c), k=n_c)

# select random samples
samples = []
for i in range(n_desired):
    idx_a = order_a[i % n_a]
    idx_b = order_b[i % n_b]
    idx_c = order_c[i % n_c]

    samples.append((idx_a, idx_b, idx_c))

(I know this code can be written a bit differently, e.g. combining all the operations on a, b, and c using list comprehensions or using itertools.cycle instead of i % n indexing, but I find this more readable, also because a, b, and c have meaningful—but irrelevant for this question—names in the original code.)

By generating random orderings of the possible values of a, b, and c and cycling through them, we ensure that the number of occurrences of the parameter values never differ by more than 1 (first, all parameter values are selected once, then twice, then thrice, etc.)

We can verify that this code achieves the desired result (equal representation (±1) of all possible parameter values):

from collections import Counter

count_a = Counter()
count_b = Counter()
count_c = Counter()

count_a.update(sample[0] for sample in samples)
count_b.update(sample[1] for sample in samples)
count_c.update(sample[2] for sample in samples)

print(f'a values are represented between {min(count_a.values())} and {max(count_a.values())} times')
print(f'b values are represented between {min(count_b.values())} and {max(count_b.values())} times')
print(f'c values are represented between {min(count_c.values())} and {max(count_c.values())} times')

This prints the following result:

a values are represented between 4 and 4 times
b values are represented between 5 and 5 times
c values are represented between 33 and 34 times

We can also verify that this code does not select duplicate combinations of a, b, and c, using the property of sets that they do not allow duplicate values:

print(len(set(samples)))

This prints 100, matching the value of n_desired.

However, an issue with this implementation is that it only works if n_desired ≤ lcm(n_a, n_b, n_c), where lcm() is the least common multiple (the smallest positive integer divisible by n_a, n_b, and n_c). In our example, lcm(n_a, n_b, n_c) = lcm(25, 20, 3) = 300. If we run the above implementation with n_desired > 300, we will see that the selected samples repeat with a period of 300. This is undesired, as this ignores 80% of the original dataset and doesn't allow us to select more than 300 unique samples.

A simple solution would be to make sure that lcm(n_a, n_b, n_c) = n_a × n_b × n_c, which would be the case if all three were prime. However, I would like this algorithm to work for any values, partly because I cannot ensure that all values are prime (for example, in my application, n_a is always the result of an integer squared).

Simply generating a list using itertools.product(range(n_a), range(n_b), range(n_c)) provides me with all possible combinations, but these are ordered sequentially, and by shuffling this full list and selecting the first n_desired samples, we lose the property of equal representation of all possible parameter values.

This is the point where I got stuck, since I am not knowledgeable enough in combinatorics to solve this issue nor do I know which terms I need to search for to find a solution. How would I solve this problem?

2 Answers

You can generate all random values of a, b and c (list of n_desired values for each of them), and then combine them into an array.

import random

# n is the maximal value to generate
# k is the number of samples, i.e. length of the resulting list
def generate(n, k):
    # the values that are evenly distributed
    l1 = list(range(n)) * (k // n)
    # remaining values that are generated one time more than another ones
    l2 = random.sample(range(n), k % n)
    l = l1 + l2
    random.shuffle(l)
    return l
    
n_a = 25
n_b = 20
n_c = 3
n_desired = 100
l = list(zip(generate(n_a, n_desired), generate(n_b, n_desired), generate(n_c, n_desired)))
print(l)

Duplicates can be removed and then resampled using this function. At first it splits the list into unique and duplicated samples. Then it tries to rearrange duplicate values so than new unique samples are generated. If it fails to decrease the number of duplicates, then it tries to remove some unique samples and use them to generate new samples.

def remove_duplicates(l):
    unique = set()
    duplicates = []
    for t in l:
        if t in unique:
            duplicates.append(t)
        else:
            unique.add(t)
    n_duplicates = len(duplicates)
    
    # iterations = 0
    # n_retries = 0
    while n_duplicates > 0:
        while n_duplicates > 0:
            # iterations += 1
            # print(n_duplicates)
            a, b, c = map(list, zip(*(duplicates)))
            for x in a, b, c:
                random.shuffle(x)
            duplicates = []
            for t in zip(a, b, c):
                if t in unique:
                    duplicates.append(t)
                else:
                    unique.add(t)
            if len(duplicates) == n_duplicates:
                break
            n_duplicates = len(duplicates)
        if n_duplicates > 0:
            # n_retries += 1
            n_recycled = min(n_duplicates, len(unique))
            recycled = random.sample(list(unique), n_recycled)
            unique = unique - set(recycled)
            duplicates += recycled
    # print(iterations, n_retries)
    return unique

It works well if n_desired is less than half of all possible samples (n_a * n_b * n_c), but takes a lot of iterations to complete otherwise. This issue can be solved by generating samples that are not included in the final set:

if n_desired <= n_a * n_b * n_c // 2:
    result = generate_samples(n_a, n_b, n_c, n_desired)
else:
    excluded = generate_samples(n_a, n_b, n_c, n_a * n_b * n_c - n_desired)
    all_samples = set(itertools.product(range(n_a), range(n_b), range(n_c)))
    result = all_samples - set(excluded)

This is much more fiddly than I was expecting, I've ended up with quite a bit of code:

from math import prod
from itertools import product
from random import shuffle

def sample(n, ns):
    # make sure parameters are valid
    if n > prod(ns):
        raise ValueError("more values requested than unique combinations", n, ns)

    # "remain" keeps track of the remaining counts for each item
    remain = []
    for n_i in ns:
        k, m = divmod(n, n_i)
        # start with the whole number
        d = {i: k for i in range(n_i)}
        # add in the remainders
        if m:
            r = list(range(n_i))
            shuffle(r)
            for i in r[:m]:
                d[i] += 1
        # sanity check
        assert(sum(d.values()) == n)

        remain.append(d)

    # generate list of all available options in random order
    opts = list(product(*(range(n_i) for n_i in ns)))
    shuffle(opts)

    result = []
    for _ in range(n):
        # get next random item, fails if we've been unlucky
        tup = opts.pop()
        result.append(tup)
        
        # keep track of remaining counts
        for i, (rem, a) in enumerate(zip(remain, tup)):
            j = rem[a]
            if j > 1:
                rem[a] = j - 1
            else:
                del rem[a]
                # remove options that involve a number that's been used up
                opts[:] = filter(lambda t: t[i] != a, opts)

    # we're done
    return result

can be used as:

x = sample(100, (25, 20, 3))

note that this starts by generating all possible options. this seems like a reasonable trade off for your parameters, but you shouldn't use this algorithm if there are billions of possible options.

also note that large ns cause this algorithm to fail, see the plot below.

success rate of n

feel free to suggest improvements, or just put it into a loop that retries on IndexError!

Related