append elements from specific positions of a list to another lists in one line function

Question:

I need to generate M lists from another list with N=a*M elements (a>=1), so the nth element of i partition is added to the the nth position of the i new list (i=1..M, so the original list is splitted in equal parts of N/M length).

With an example is a lot easier to understand:
so supose that I have this list lg=[1,2,3,4,5,6,7,8,9], the result should be l1=[1,4,7], l2=[2,5,8], l3=[3,6,9]

I did it with this code:

lg=[1,2,3,4,5,6,7,8,9] # N=9
M=3
lr=[[] for _ in range(M)]

print(lr) # [[], [], []]
idx_lr = 0
for idx,el in enumerate(lg): # the enumerate is not necessary, was for testing purposes
  lr[idx_lr].append(el)
  idx_lr += 1
  if idx_lr == M:
    idx_lr = 0
print(lr) #[[1, 4, 7], [2, 5, 8], [3, 6, 9]]

and it works, but, it’s an ugly code, there is a way to do it with one line function?

Asked By: NewbieInFlask

||

Answers:

You can do this with a simple slice operation:

>>> lg=[1,2,3,4,5,6,7,8,9]
>>> [lg[i::3] for i in range(3)]
[[1, 4, 7], [2, 5, 8], [3, 6, 9]]

The third part of the slice is the "step" between each element in the slice; it defaults to 1, but setting it to M makes it take every Mth element.

Answered By: Samwise
Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.