CODEWITHSUNDEEP
pythonlistrandom
base12masterrace
base12masterrace

Reputation: 27

Randomly sample from arbitrarily nested lists while maintaining structure

I am trying to write a function which chooses integers randomly from arbitrarily nested lists while maintaining order and the structure of the lists (though empty lists can be ignored). The number of samples should be uniformly random from 0 to the number of integers in the nested list.

For example, running listrand([1,2,3,[3,4],[65,[3]]]) 3 times might give:

[1, 3, [3], [65, [3]]]
[1, 2, [[3]]]
[[3, 4], [65]]

The catch is I need it to be uniformly distributed, so I can't use something like

sample = [[random.sample(mylist)] for mylist in listoflists]

because that would be binomial.

At a minimum I need this to work for single-level nesting. I thought about sampling from a flattened list but then I'm not sure how to use those to construct the desired output.

Upvotes: 1

Views: 278

Answers (1)

chthonicdaemon
chthonicdaemon

Reputation: 19770

This solution satisfies your requirements by construction. In other words, the number of elements chosen is uniformly random.

import random
from collections import Iterable

def count(l, r=0):
    for i in l:
        if isinstance(i, Iterable):
            r += count(i)
        else:
            r += 1
    return r

def listrand(target):
    N = count(target)
    nchosen = random.randint(0, N-1)
    chosen = set(random.sample(range(N), nchosen))

    def build(l, c=0):
        output = []
        for i in l:
            if isinstance(i, Iterable):
                c, sublist = build(i, c)
                if sublist:
                    output.append(sublist)
            else:
                if c in chosen:
                    output.append(i)
                c += 1
        return c, output

    return build(target)[1]

Sample output:

target = [1,2,3,[3,4], [65,[3]]]

for i in range(5):
    print(listrand(target))

[1, 2, 3, [3, 4], [[3]]]
[2]
[]
[2, 3, [3, 4]]
[[4]]

Upvotes: 3

Related Questions