Most efficient way to convert a list of counts to a list of numbers

user_1726404 0
Viewed 117

I have a list of counts where each index represents a number and its count represents how many of that number is in the list:

a = [3,5,1,2]

turns into

b = [0,0,0,1,1,1,1,1,2,3,3]

I was thinking we could do something like:

b = []
for ix, el in enumerate(a):
    b.extend([ix]*a[ix])
print(b)

But if I am not mistaken it takes k (count val) time to put it in list b as the extend takes k time, but it also has to be done n times giving us a runtime of n*k where n is the number of indicies and k is the number of counts for each index

Another idea is instead of having an array of counts we can have an array of the pure elements:

a = [[0,0,0],[1,1,1,1,1],[2],[3,3]]

but to flatten it it still takes quite some time (I believe n*k time)

b = [item for sublist in a for item in sublist]

is there a way to make this more efficient? Maybe converting to string removing all brackets and converting back into a list?

time-complexity python
1 Answers

You can use numpy's np.repeat for a performant approach:

np.repeat(np.arange(len(a)), a)
# array([0, 0, 0, 1, 1, 1, 1, 1, 2, 3, 3])

Here are the timings -

a_large = np.concatenate([a]*10_000, axis=0)

def op(a):
    b = []
    for ix, el in enumerate(a):
        b.extend([ix]*a[ix])

def yatu(a):
    np.repeat(np.arange(len(a)), a)

%timeit op(a_large)
# 17.1 ms ± 422 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit yatu(a_large)
# 368 µs ± 1.91 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
user_96986840
Related