CODEWITHSUNDEEP
pythonlistlist-comprehension
Jack Rolph
Jack Rolph

Reputation: 595

Using List Comprehension To Repeat Entries in 2D Array

I have the following piece of code, with the goal of repeating every element of a 2D list a certain number of times:

empty =[]
l = [[1,2,3], [4,5,6]]
n = [2,3]
for i, _as in enumerate(l):
    for _a in _as:
        for _n in range(n[i]):
            empty.append(_a)

empty
> [1, 1, 2, 2, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6, 6]

I would like to try and write this in a single line list comprehension format.

I have made an attempt using using:

empty = [
         [
           [_a]*n[i] for _a in _as
         ] 
         for i, _as in enumerate(l) 
        ]

empty
> [[[1, 1], [2, 2], [3, 3]], [[4, 4, 4], [5, 5, 5], [6, 6, 6]]]

How may I correct the above code to provide the desired result?

Upvotes: 0

Views: 86

Answers (3)

RoadRunner
RoadRunner

Reputation: 26315

You could use itertools.repeat to repeat the elements of each sublist with n, then itertools.chain.from_iterable to flatten the final list:

>>> from itertools import chain, repeat
>>> l = [[1,2,3], [4,5,6]]
>>> ns = [2, 3]
>>> list(chain.from_iterable(repeat(x, n) for lst, n in zip(l, ns) for x in lst))
[1, 1, 2, 2, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6, 6]

Upvotes: 0

Red
Red

Reputation: 27557

You can use a double list comprehension:

l = [[1,2,3], [4,5,6]]
n = [2, 3]
empty =[i for j in [[a]*f for b,f in zip(l,n) for a in b] for i in j]
print(empty)

Output:

[1, 1, 2, 2, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6, 6]

Upvotes: 0

Vishal Singh
Vishal Singh

Reputation: 6234

if you're just looking to convert your code into a list comprehension then just write your loops in their original order. This way you'll never forget how to get things done using list comprehension.

output = [
    _a
    for i, _as in enumerate(l)
    for _a in _as
    for _n in range(n[i])
]

Upvotes: 5

Related Questions