Python中使用多个列表进行数学计算

Question

我正在尝试在Python中获取三个不同列表的所有可能组合。

假设我有三个列表。

A = [2,3,4,7,9]
B = [2,3,5]
C = [1,2]

我想返回所有满足 A[i]+B[j]-C[k] == 3 的组合。

D = [[i,j,k] for i in A for j in B for k in C]

这个句子的翻译是：“这会给我一个列表格式的所有组合，但接下来该怎么做呢？”

- JesterB

Answer 1

您可以使用带有itertools.product的列表推导式来检查每个组合，以查看它是否符合您的条件。

>>> from itertools import product
>>> [(a,b,c) for a,b,c in product(A,B,C) if a + b - c == 3]
[(2, 2, 1), (2, 3, 2), (3, 2, 2)]

- Cory Kramer

---

虽然我不确定这是否是 OP 的要求，但也许你应该将 (a,b,c) 改为 [a,b,c]？ - Iron Fist

Answer 2

只需添加if i+j-k == 3。不需要使用A[i]等，因为这会索引每个list，而i、j和k已经是元素本身。

D = [[i,j,k] for i in A for j in B for k in C if i+j-k == 3]

你还应该了解 itertools.product，以避免像这样嵌套的情况：

import itertools
D = [[i,j,k] for i,j,k in itertools.product(A, B, C) if i+j-k == 3]

- TigerhawkT3

---

是的，我现在明白了。我将在一两周内换上新眼镜，抱歉。 - TigerhawkT3

我不知道哪个更快，但使用 itertools 的方法更容易用于 n 个列表的编写，特别是如果你把它们都放入一个容器中，此时你只需执行 itertools.product(*container)。 - TigerhawkT3

Answer 3

只需在最后添加条件即可：

>>> D = [[i,j,k] for i in A for j in B for k in C if i + j - k == 3]
>>> D
[[2, 2, 1], [2, 3, 2], [3, 2, 2]]

- Mike Müller