每个元素可以是“正”或“负”的幂集

Question

我在寻找一种简单的方法来生成一个可迭代对象的所有子集，其中每个元素可以是“正”的或“负”的，但在同一个组合中不能同时出现这两种情况。可迭代对象中没有重复元素，只有元素本身或它的负值。顺序不重要。

下面是一个包含 int 的列表的例子：

可迭代对象：

elements = [-2, 1]

想要的所有子集：

[]
[-2]
[2]
[-1]
[1]
[-2, -1]
[-2, 1]
[2, -1]
[2, 1]

需要排除的“子集”：

[-1, 1]
[-2, 2]

我现在的方法是使用 这里 的子集实现，结合 elements 列表和 [-x for x in elements]，然后遍历所有子集并删除不需要的组合。不过，我觉得这样不是最优的。有没有简单的解决方案，不需要我最后再去删除不需要的组合呢？

Answer 1

逐个扩展幂集的功能。

elements = [-2, 1]

powerset = [[]]
for e in elements:
    powerset += [
        p + [x]
        for p in powerset
        for x in (e, -e)
    ]

print(*powerset)

输出结果（在线尝试！）：

[] [-2] [2] [1] [-1] [-2, 1] [-2, -1] [2, 1] [2, -1]

---

生成器版本：

elements = [-2, 1]

powerset = iter([[]])
for e in elements:
    powerset = (
        x
        for powerset, e in [(powerset, e)]
        for p in powerset
        for x in (p, p + [e], p + [-e])
    )

print(*powerset)

输出结果（在线尝试！）：

[] [1] [-1] [-2] [-2, 1] [-2, -1] [2] [2, 1] [2, -1]

---

再来一个，应该更快：

from functools import reduce

elements = [-2, 1]

def empower(powerset, element):
    pos = [element]
    neg = [-element]
    for subset in powerset:
        yield subset
        yield subset + pos
        yield subset + neg

powerset = reduce(empower, elements, iter([[]]))

print(*powerset)

输出结果（在线尝试！）：

[] [1] [-1] [-2] [-2, 1] [-2, -1] [2] [2, 1] [2, -1]

---

基准测试

创建和遍历幂集的时间，针对 elements = range(1, 13)：

 557 ms  no_comment_1
 554 ms  no_comment_1b
 472 ms  no_comment_1c
 165 ms  no_comment_2
  77 ms  no_comment_3
 696 ms  user2357112

基准测试脚本：

from functools import reduce
from collections import deque
from time import time
import itertools

def no_comment_1(elements):
    powerset = [[]]
    for e in elements:
        powerset += [
            p + [x]
            for p in powerset
            for x in (e, -e)
        ]
    return powerset

def no_comment_1b(elements):
    powerset = [[]]
    for e in elements:
        xs = [e], [-e]
        powerset += [
            p + x
            for p in powerset
            for x in xs
        ]
    return powerset

def no_comment_1c(elements):
    powerset = [[]]
    for e in elements:
        x = [e]
        pos = [p + x for p in powerset]
        x = [-e]
        neg = [p + x for p in powerset]
        powerset += pos
        powerset += neg
    return powerset

def no_comment_2(elements):
    powerset = iter([[]])
    for e in elements:
        powerset = (
            x
            for powerset, e in [(powerset, e)]
            for p in powerset
            for x in (p, p + [e], p + [-e])
        )
    return powerset

def no_comment_3(elements):
    def empower(powerset, element):
        pos = [element]
        neg = [-element]
        for subset in powerset:
            yield subset
            yield subset + pos
            yield subset + neg
    return reduce(empower, elements, iter([[]]))

def user2357112(elements):
    elements = list(elements)
    n = len(elements)
    for option in itertools.product([-1, 0, 1], repeat=n):
        yield [i*elem for (i, elem) in zip(option, elements) if i]

funcs = no_comment_1, no_comment_1b, no_comment_1c, no_comment_2, no_comment_3, user2357112

elements = [-2, 1]
for f in funcs:
    print(*sorted(f(elements)))

elements = range(1, 13)
times = {f: [] for f in funcs}
for _ in range(5):
    for f in funcs:
        t = time()
        deque(f(elements), 0)
        times[f].append(time() - t)
for f in funcs:
    t = min(times[f])
    print(f'{t*1e3:4.0f} ms ', f.__name__)

在线尝试！

Answer 2

每个元素可以被包含、被排除或者被否定。可以使用 itertools.product 来遍历所有 3**len(elements) 种可能的选择方式，也就是对每个元素进行选择的所有组合：

import itertools

def extended_powerset(elements):
    elements = list(elements)
    n = len(elements)
    for option in itertools.product([-1, 0, 1], repeat=n):
        yield [i*elem for (i, elem) in zip(option, elements) if i]