как изменить значения списка на основе индекса другого списка в python?

Question

У меня есть следующий список

I=[[[0, 0, 0], [0, 0, 0], [0, 0, 0]],
   [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
   [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
   [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
   [[0, 0, 0], [0, 0, 0], [0, 0, 0]]]

Я хотел бы заменить некоторые элементы от 0 на 1 в списке I на основе индекса другого списка S

S=[2,2,1]

Обратите внимание, что: Первый элемент в списке S соответствует первому столбцу в списке I, 2-й элемент в списке S соответствует 2-му столбцу в списке I 3-й элемент в списке S соответствует 3-му столбцу в списке I

вывод, подобный этому

IS=[[[1, 0, 0], [0, 0, 0], [0, 0, 0]],  
    [[0, 1, 0], [0, 0, 0], [0, 0, 0]],     
    [[0, 0, 0], [1, 0, 0], [0, 0, 0]],
    [[0, 0, 0], [0, 1, 0], [0, 0, 0]],     
    [[0, 0, 0], [0, 0, 0], [1, 0, 0]]]

Вывод, который я хочу получить, выглядит следующим образом: начиная с colmn1, который соответствует первому элементу в списке S, поскольку S [0] = 2, мы получаем доступ к первым двум строкам в столбце 1 списка я и заменяю 0 на 1 по диагонали. тогда мы получим

 colmn1      colmn2      colmn3
[[**1, 0, 0]**, [0, 0, 0], [0, 0, 0]], row 1 
[**[0, 1, 0]**, [0, 0, 0], [0, 0, 0]], row 2

Затем перейдем ко второму элементу S, который соответствует colmn2 в списке I, так как S [1] = 2, тогда мы получим доступ к следующим двум строкам в colmn2 и заменим 0 элементов по диагонали в следующих двух строках.

 colmn1      colmn2      colmn3
[[1, 0, 0], [0, 0, 0], [0, 0, 0]], row 1 
[[0, 1, 0], [0, 0, 0], [0, 0, 0]], row 2
[[0, 0, 0], **[1, 0, 0]**, [0, 0, 0]], row 3
[[0, 0, 0], **[0, 1, 0]**, [0, 0, 0]], row 4

наконец, go до последнего элемента в S (соответствующего последнего столбца в списке I), так как мы имеем S [2] = 1, так что мы получаем доступ к следующей строке и заменить 0 на 1 в диагональном направлении, чтобы получить

 colmn1      colmn2      colmn3
[[1, 0, 0], [0, 0, 0], [0, 0, 0]], row 1 
[[0, 1, 0], [0, 0, 0], [0, 0, 0]], row 2
[[0, 0, 0], [1, 0, 0], [0, 0, 0]], row 3
[[0, 0, 0], [0, 1, 0], [0, 0, 0]], row 4 
[[0, 0, 0], [0, 0, 0], **[1, 0, 0]**], row 5

Я надеюсь, что объяснение понятно

Answer 1

вы можете использовать:

cols = zip(*I)

offset = 0
for l, n in zip(cols, S):
    for i in range(n):
        l[offset][i] = 1
        offset += 1

print(I)

вывод:

[[[1, 0, 0], [0, 0, 0], [0, 0, 0]],
 [[0, 1, 0], [0, 0, 0], [0, 0, 0]],
 [[0, 0, 0], [1, 0, 0], [0, 0, 0]],
 [[0, 0, 0], [0, 1, 0], [0, 0, 0]],
 [[0, 0, 0], [0, 0, 0], [1, 0, 0]]]