Python map a value to each i-th sublist's element

Question

I'm trying to do the following in python: given a list of lists and an integer i

input = [[1, 2, 3, 4], [1, 2, 3, 4], [5, 6, 7, 8]]
i = 1

I need to obtain another list which has all 1s for the elements of the i-th list, 0 otherwise

output = [0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0]

I wrote this code

output = []
for sublist in range(0, len(input)):
    for item in range(0, len(input[sublist])):
        output.append(1 if sublist == i else 0)

and it obviously works, but since I'm a newbie in python I suppose there's a better 'pythonic' way of doing this.

I thought using map could work, but I can't get the index of the list with it.

Answer 1

Creating extra variable to get index of current element in interation is quite unpythonic. Usual alternative is usage of enumerate built-in function.

Return an enumerate object. sequence must be a sequence, an iterator, or some other object which supports iteration. The next() method of the iterator returned by enumerate() returns a tuple containing a count (from start which defaults to 0) and the values obtained from iterating over sequence.

You may use list comprehension with double loop inside it for concise one liner:

input_seq = [[1, 2, 3, 4], [1, 2, 3, 4], [5, 6, 7, 8]]
i = 1
o = [1 if idx == i else 0 for idx, l in enumerate(input_seq) for _ in l]

Alternatively,

o = [int(idx == i) for idx, l in enumerate(input_seq) for _ in l]

Underscore is just throwaway name, since in this case we don't care for actual values stored in input sublists.

Answer 2

Here's a 1-liner, but it's not really obvious:

output = [int(j == i) for j, sublist in enumerate(input) for _ in sublist]

Somewhat more obvious:

output = []
for j, sublist in enumerate(input):
    output.extend([int(i == j)] * len(sublist))

Then "0 or 1?" is computed only once per sublist, which may or may not be more efficient.