如何遍历嵌套列表的一部分？

Question

I have this nested list:

nested = [[1, 1, 1], [2, 2, 2], [3, 3, 3]]

I wanted to subtract 1 from the first 2 elements of the inner list so I tried it with list comprehension:

nested = [[x - 1 for x in stack[0:2]] for stack in nested]

It did give me back the first 2 elements subtracted by 1 for the inner lists but it removed the last element completely

nested = [[0, 0], [1, 1], [2, 2]]

I thought that by slicing the list, it will not affect the other element. However in this case it didn't work. Can someone explain this to me?

Answer 1

To keep the 3rd element, include it in the list comprehension:

>>> [ [x - 1 for x in stack[0:2]] + stack[2:] for stack in nested ]
[[0, 0, 1], [1, 1, 2], [2, 2, 3]]

The above works for stack of any length.

Or, if stack always has exactly three elements:

>>> [[x-1, y-1, z] for x, y, z in nested]
[[0, 0, 1], [1, 1, 2], [2, 2, 3]]

Or, you can make the changes in place:

>>> for stack in nested: stack[0]-=1; stack[1]-=1
... 
>>> nested
[[0, 0, 1], [1, 1, 2], [2, 2, 3]]

Answer 2

Another option is to use numpy which does this sort of slicing naturally

import numpy as np
nested = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])

nested[:,:2] -= 1

returns

array([[0, 0, 1],
       [1, 1, 2],
       [2, 2, 3]])

Answer 3

Try it like this:

nested = [[x - 1  if i < 2 else x for i,x in enumerate(stack)] for stack in nested]

This will affect only first two elements keeping the rest as is.

Answer 4

You can use enumerate :

nested = [[1, 1, 1], [2, 2, 2], [3, 3, 3]]
new_nested = [[a-1 if i == 0 or i == 1 else a for i, a in enumerate(b)] for b in nested]

Output:

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

Edit: alternative, with map :

nested = [[1, 1, 1], [2, 2, 2], [3, 3, 3]]
new_nested = [map(lambda x:x-1, i[:2])+[i[-1]] for i in nested]

Output:

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