Repeat numbers using itertools

user_5200329 0
Viewed 292

I was learning itertools and encountered an interesting problem.

How to get the following result?

nums = [1,2,3,4]
# logic:  i+1 th number is repeated i times.
          2 is repeated 1 times and 4 is repeated 3 times.

required = [2,4,4,4]

My attempt

import itertools

nums = [1,2,3,4]
nums[::2]  # [1,3]
nums[1::2] # [2,4]
              * 2 is repeated 1 times
                * 4 is repeated 3 times and makes [2,4,4,4]

list(itertools.starmap(itertools.repeat, zip(nums[::2],nums[1::2])))

Gives,
[repeat(1, 2), repeat(3, 4)]

How to get:
required = [2,4,4,4]

Required

From: [1,2,3,4]
To:   [2,4,4,4]

We can use itertools, list comp and numpy.
With emphasis on itertools, since I am learning itertools functions
such as repeat, starmap and so on.
itertools python
3 Answers

repeat creates an iterable object, so you need to chain those iterators together, and before creating a list by iterating over the result. (You also need to swap the arguments to zip)

>>> from itertools import starmap, repeat, chain
>>> list(chain.from_iterable(starmap(repeat, zip(nums[1::2], nums[::2]))))
[2, 4, 4, 4]

Actually, you don't need either starmap or zip; map can apply a multi parameter function to multiple iterators, taking one argument from each iterator per call.

>>> list(chain.from_iterable(map(repeat, nums[1::2], nums[::2])))
[2, 4, 4, 4]

Going one step further into itertools land, you can use tee and islice to avoid having to create two temporary lists by slicing nums (though doing so in one line gets messy):

>>> from itertools import tee, islice
>>> t1, t2 = tee(nums)
>>> list(chain.from_iterable(map(repeat, islice(t1, 1, None, 2), islice(t2, None, None, 2))))
[2, 4, 4, 4]
user_11268410

you don't need itertools for that. A list comprehension and zip are sufficient:

nums = [1,2,3,4]

rnums = [ rn for r,n in zip(nums[::2],nums[1::2]) for rn in r*[n] ]

print(rnums) # [2, 4, 4, 4]

[EDIT]: using itertools:

nums = [*chain(*map(repeat,zip(nums[::2],nums[1::2])))]
user_52375600

If you want to go simple, no itertools, no list comprehension, it can be done very easily like this:

nums = [1,2,3,4]
ans = []
for i in range(len(nums)):
    if i%2: # odd index
        ans += [ nums[i] ] * nums[i-1]

Using list comp

sum(([nums[i+1]]*nums[i] for i in range(0, len(nums), 2)), [])
user_52003290
Related