Home > Back-end >  Set desired element to 1, the rest 0
Set desired element to 1, the rest 0

Time:03-17

My goal is to set only the desired elements of the my_list to 1 and the rest to 0. Expected output is [1,0,1]. Is there a simply way without using explicit for loop?

my_list = [0.2, -0.5, 0.1]
set_to_one = [0, 2]  # I want to set 0 th, 2nd element to 1, the rest 0

Expected result is [1,0,1].

CodePudding user response:

Perhaps build a list of zeros, then run a loop setting the desired elements to one.

my_list = [0] * len(my_list)
for i in set_to_one:
    my_list[i] = 1

Result:

>>> my_list
[1, 0, 1]

CodePudding user response:

Using starmap from itertools you can unpack the tuples generated by enumerate and then classify.

import itertools as it

to_ones = [0, 2]

my_list = [0.2, -0.5, 0.1]

iter_ = it.starmap(lambda i, _:1 if i in to_ones else 0, enumerate(my_list))

print(list(iter_))

Remark: since the output list does not depends on its elements but only on its length, the iterator passed to starmap is a bit an "overkill". [I readapted my original approach which used a list comprehension]

  • Related