Powerset where each element can be either "positive" or "negative"

Question

I'm looking for a simple way to generate the powerset of an iterable where each element can be "positive" or "negative", but not both in the same combination. There are no duplicates in the iterable and only the element or it's negative. Order doesn't matter.

Here is an example with a list of int's:

The iterable:

elements = [-2, 1]

The desired powerset:

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

"Subsets" to exclude:

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

My current approach is to use the powerset implementation from here for the combined list of elements and [-x for x in elements] and then iterate through the powerset and remove unwanted combinations. However, that's not optimal I guess. Is there a simple solution that doesn't require me to remove unwanted combinations in the end?

Answer 1

Each element can be included, included negated, or excluded. Use itertools.product to go through all 3**len(elements) possible ways to pick an option for each element:

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]

Answer 2

Extending the power set one element at a time.

elements = [-2, 1]

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

print(*powerset)

Output (Attempt This Online!):

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

---

A generator version:

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)

Output (Attempt This Online!):

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

---

One more, should be faster:

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)

Output (Attempt This Online!):

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

---

Benchmark

Times for creating and iterating the powerset for 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

Benchmark script:

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__)

Attempt This Online!