CODEWITHSUNDEEP
pythonnumpy
Kevin Yang
Kevin Yang

Reputation: 167

Python iterator that generates combinatorial splits

I know that itertools.combinations will iterate through all combinations of a certain length from an iterable:

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)

Would print 

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

However, I want something like 

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

To print 

(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])

Upvotes: 0

Views: 164

Answers (2)

wim
wim

Reputation: 362796

I would use combinations to generate a mask instead, and then just use the mask to fancy-index your array. That would look like this:

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]

Upvotes: 1

user2357112
user2357112

Reputation: 280993

This doesn't come with Python, but you can build it on top of itertools.combinations fairly easily:

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

Upvotes: 0

Related Questions