Search code examples
pythonarrayspython-3.xnumpypython-itertools

Python: Finding all possible unique two-index combinations in a 2D array


I have a numpy array that is roughly equivalent to:

data = ([1, 2, 3], [4, 5, 6], [7, 8, 9])

I want to find all the unique two-index combinations of these values. In other words, I want all the possible combinations without repeating a row or column index (similar to correctly solving a Sudoku puzzle). For example, the desired output would be:

output >> ([1, 5, 9], [1, 6, 8], [2, 4, 9], [2, 6, 7], [3, 4, 8], [3, 5, 7])

and this output can be represented by their corresponding indices: ([0][0],[1][1],[2][2]), ([0][0],[1][2],[2][1]), ([0][1],[1][0],[2][2]), ([0][1],[1][2],[2][0]), ([0][2],[1][0],[2][1]), ([0][2],[1][1],[2][0])

I've tried using itertools.permutations, and while it does find all the possible permutations of my data for each unique row, it does not treat each column as unique)

I want only one value from each row and each column.

I’m fairly new to python, does anyone have a suggestion of how I might do this?


Solution

  • from itertools import permutations
    
    data = ([1, 2, 3], [4, 5, 6], [7, 8, 9])
    
    output = [[row[y] for row, y in zip(data, permutation)]
              for permutation in permutations(range(len(data)))]
    

    EDIT: The problem has changed in the comments to only yield results that don't contain 0. Also since len(data) is 100 we can't produce all results using permutations like above and then filter them; that would take forever. They have to be correctly selected as we go, like so:

    def get_nonzero_perms(data, remaining_indices=None):
        if not data:
            yield []
            return
        if remaining_indices is None:
            remaining_indices = list(range(len(data)))
        row = data[0]
        for i in remaining_indices:
            value = row[i]
            if value == 0:
                continue
            for perm in get_nonzero_perms(data[1:], [j for j in remaining_indices if j != i]):
                yield [value] + perm
    
    
    for p in get_nonzero_perms(([2, 8, 0, 0], [0, 3, 9, 4], [0, 0, 5, 1], [4, 6, 0, 7])):
        print(p)
    

    Output:

    [2, 3, 5, 7]
    [2, 9, 1, 6]
    [2, 4, 5, 6]
    [8, 9, 1, 4]
    [8, 4, 5, 4]