List comprehension inside a list comprehension with comparison condition

Question

I have a list like this,

 sl=[[1,2,100],[2,100,4],[100,4,5],[4,5,6],[5,6,200],[6,200,7],[200,7,8],[7,8,300],[8,300,9]]

Now I want to to find those element where the element is greater than the mean and store it into a list.

so, the list will look like,

  [[100],[100],[100],[],[200],[200],[200],[300]]

I can do it using for loop, the code is following,

indices=[]
for i in sl:
    indices.append([j for j in i if (j>(np.mean(i))])

But the execution time is long to execute. I want to avid the for loop and use some kind of list comprehension to do the same task.

Is there any way to do it most efficient execution time?

Answer 1

You should calculate the mean only once per loop. This should improve things a bit:

indices = []
for i in sl:
  mean = np.mean(i)
  indices.append([j for j in i if j > mean])

I don't think converting the whole thing to a nested list comprehension would help, because we want to extract the mean calculation outside of the innermost loop.

Answer 2

For a regular shaped list, ie same number of elements per list at the inner nested level, we can use NumPy tools to offload the compute part, like so -

a = np.array(sl)
m = a>a.mean(1,keepdims=True)
idx = np.r_[0,m.sum(1).cumsum()]
f = a[m].tolist()
out = [f[i:j] for (i,j)in zip(idx[:-1],idx[1:])]

This would be beneficial for a large number of entries in sl , ie for a large len(sl) .

Answer 3

you can have everything in one list comprehension avoiding for loop, getting a bit of performance:

means = map(np.mean, sl)
indices = [list(filter(lambda x: x > m, l)) for l, m in zip(sl, means)]

Answer 4

You can also do:

sl=[[1,2,100],[2,100,4],[100,4,5],[4,5,6],[5,6,200],[6,200,7],[200,7,8],[7,8,300],[8,300,9]]

sl = np.array(sl)

# List of elements satisfying the condition
elem_list = [sl[i][m].tolist() for i,m in enumerate((sl.T>sl.mean(axis=1)).T)]

# List of index for which the respective elements satisfy the condition
index_list = [np.where(m)[0].tolist() for m in (sl.T>sl.mean(axis=1)).T]

Which results in:

elem_list
[[100], [100], [100], [6], [200], [200], [200], [300], [300]]

index_list
[[2], [1], [0], [2], [2], [1], [0], [2], [1]]