如何一次将两个项目添加到列表理解中

Question

我有一个简单的 for 循环，如下所示：

for i in range(3):
    if i == 1:
        for x in [0, 2]:
            free_spots.append([x, i])
    else:
        free_spots.append([1, i])

输出这个列表[[1, 0], [0, 1], [2, 1], [1, 2]]

我计划将其转换为这样的列表理解，但遇到了一个问题：

[[x, i] for x in [0, 2] if i == 1 else [1, i] for i in range(3)]

这会导致错误，我已经尝试像这样修复它：

[[[x, i] for x in [0, 2]] if i == 1 else [1, i] for i in range(3)]

虽然现在运行它也会创建这个： [[1, 0], [[0, 1], [2, 1]], [1, 2]]并且我不希望第 2 项和第 3 项成为在一个列表中。 是否可以这样做？： [[0,i],[2,i] if i == 1 else [1, i] for i in range(3)]

Answer 1

如果你真的想在一行中完成，这个怎么样？

free_spots = [val for sublist in [([[x, i] for x in [0,2]] if i == 1 else [[1, i]]) for i in range(3)] for val in sublist]

# This effectively creates a list of lists containing the values you want
main_list = [([[x, i] for x in [0,2]] if i == 1 else [[1, i]]) for i in range(3)]
# And then flattens it
free_spots = [val for sublist in main_list for val in sublist]

但其他人是对的。 除非您有充分的理由，否则这不是一个好主意。 我几乎无法理解我自己的代码，因为我已经编写了它不是 python zen。

Answer 2

你可以分两步完成。

free_spots = [[1, i] for i in range(3)]
free_spots[1:2] = [[x, 1] for x in [0, 2]]

切片赋值相当于for循环中的嵌套循环。

但是，如果循环中有多个异常，则这不能很好地概括。

Answer 3

正如评论中已经指出的那样，创建这样的列表理解语句并不优雅。 因此，正如Karl's comment所建议的，生成器是首选解决方案。

要使用生成器，您的代码将以某种方式如下所示：

# Create a generator function
def generate_data():
    for i in range(3):
        if i == 1:
            for x in [0, 2]:
                yield [x, i]
        else:
            yield [1, i]

# Use it
print(list(generate_data()))