Python 3.6 Create mapping from an array with sizes

Question

Task

I have an array with sizes

sizes = [5, 3, 1, 2]

and based on the sizes I want to create the following array

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

Solution

My first attempt

mapping = []
ctr = 0
for i in range(len(sizes)):
    for j in range(sizes[i]):
        mapping.append(i)
        ctr += 1

Shorter version

mapping = [[i for _ in range(sizes[i])] for i in range(len(sizes))]
mapping = list(itertools.chain(*mapping))

Question

One line version?

Is it possible to do it in just one line with a neat code?

Rakesh · Accepted Answer · 2020-03-19 13:11:40Z

2

Using enumerate

Ex:

sizes = [5, 3, 1, 2]
result = [i for i, v in enumerate(sizes) for _ in range(v)]
print(result)

Output:

[0, 0, 0, 0, 0, 1, 1, 1, 2, 3, 3]

answered Mar 19, 2020 at 13:11

Rakesh

82.9k17 gold badges85 silver badges122 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

RoadRunner · Accepted Answer · 2020-03-19 13:41:36Z

1

Another approach would be to multiply the indices into sublists [[0, 0, 0, 0, 0], [1, 1, 1], [2], [3, 3]] then flatten the result with itertoo.chain.from_iterable:

>>> from itertools import chain
>>> sizes = [5, 3, 1, 2]
>>> list(chain.from_iterable([i] * x for i, x in enumerate(sizes)))
[0, 0, 0, 0, 0, 1, 1, 1, 2, 3, 3]

answered Mar 19, 2020 at 13:41

RoadRunner

26.4k6 gold badges46 silver badges77 bronze badges

Collectives™ on Stack Overflow

Python 3.6 Create mapping from an array with sizes

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related