Создайте больший список из существующего списка, используя понимание списка или карту ()

Question

Я пытаюсь создать список для индексации координат (x, y и z), учитывая набор индексов атомов. Моя проблема состоит в том, как просто изящно выйти из этого списка:

atom_indices = [0, 4, 5, 8]

К этому списку:

coord_indices = [0, 1, 2, 12, 13, 14, 15, 16, 17, 24, 25, 26]

Самым простым способом чтения / понимания, о котором я думал до сих пор, является простой:

coord_indices = []
for atom in atom_indices:
    coord_indices += [3 * atom,
                      3 * atom + 1,
                      3 * atom + 2]

Но это не похоже на Pythonic. Есть ли лучший способ, о котором я не думал, не получив список списков или список кортежей?

Answer 1

Как насчет:

>>> atom_indices = [0, 4, 5, 8]
>>> coords = [3*a+k for a in atom_indices for k in range(3)]
>>> coords
[0, 1, 2, 12, 13, 14, 15, 16, 17, 24, 25, 26]

Мы можем вкладывать циклы в списки в том же порядке, в котором мы пишем циклы, т.е.

coords = []
for a in atom_indices: 
    for k in range(3): 
        coords.append(3*a+k)

Не бойся for петли, хотя, если они более понятны в ситуации. По причинам, которые я никогда полностью не понимал, некоторые люди чувствуют себя более умными, когда пишут код горизонтально, а не вертикально, хотя это затрудняет отладку.