Question

我知道itertools.combinations会迭代迭代中某个长度的所有组合：

import itertools
import numpy as np
my_array = np.array([[0,1,2],[3,4,5],[6,7,8]])
for c in itertools.combinations(my_array, 2):
    print (c)

会打印

(array([0, 1, 2]), array([3, 4, 5]))
(array([0, 1, 2]), array([6, 7, 8]))
(array([3, 4, 5]), array([6, 7, 8]))

但是，我想要像

这样的东西

for c, d in combo_splits(my_array, 2):
    print (c)
    print (d)

打印

(array([0, 1, 2]), array([3, 4, 5])) 
array([4, 5, 6])
(array([0, 1, 2]), array([6, 7, 8])) 
array([3, 4, 5])
(array([3, 4, 5]), array([6, 7, 8]))
array([0, 2, 3])

Answer 1

这不是Python附带的，但你可以很容易地在itertools.combinations之上构建它：

def combinations_with_leftovers(pool, k):
    """itertools.combinations, but also returning the parts we didn't take.

    Each element of combinations_with_leftovers(pool, k) is a tuple (a, b)
    where a is the tuple that itertools.combinations would have given
    and b is a tuple of the elements of pool not used in a.

    """
    pool = tuple(pool)
    for chosen_indices in itertools.combinations(xrange(len(pool)), k):
        chosen_indices = set(chosen_indices)
        a = tuple(x for i, x in enumerate(pool) if i in chosen_indices)
        b = tuple(x for i, x in enumerate(pool) if i not in chosen_indices)
        yield a, b

Answer 2

我会使用组合来生成掩码，然后使用掩码来对数组进行花式索引。这看起来像这样：

import itertools
import numpy as np

my_array = np.array([[0,1,2],[3,4,5],[6,7,8]])

n = len(my_array)

for r in itertools.combinations(xrange(n), 2):
    rows = np.array(r)
    on = np.zeros(n, dtype=bool)
    on[rows] = True
    print my_array[on]
    print my_array[~on]

生成组合拆分的Python迭代器

2 个答案: