CODEWITHSUNDEEP
pythonnumpycombinatorics
ZMP
ZMP

Reputation: 47

How do I make copies of an array in python based on conditions in the array?

I have an 

n = np.array([[1, 12, 1, 3],
              [1, 1, 12, 0]])

and would like to duplicate it such that if I have a double-digit number in the array, it breaks the array into two identical arrays where the first array has the first digit and the second array has the second digit. In the above example, I would have 4 copies of the matrix. The assumptions are that there are either single digit or double digit numbers in the array.

n1 = [1, 1, 1, 3], [1, 1, 1, 0]
n2 = [1, 1, 1, 3], [1, 1, 2, 0]
n3 = [1, 2, 1, 3], [1, 1, 1, 0]
n4 = [1, 2, 1, 3], [1, 1, 2, 0]

Upvotes: 1

Views: 115

Answers (1)

Paul Panzer
Paul Panzer

Reputation: 53029

Approach 1: itertools.product

>>> import numpy as np
>>> from itertools import product
>>> from pprint import pprint
>>> 
>>> n = np.array([[1, 12, 1, 3],
...               [1, 1, 12, 0]])
>>> 
>>> pprint([np.reshape(nn, n.shape).astype(int) for nn in product(*map(str, n.ravel()))])
[array([[1, 1, 1, 3],
       [1, 1, 1, 0]]),
 array([[1, 1, 1, 3],
       [1, 1, 2, 0]]),
 array([[1, 2, 1, 3],
       [1, 1, 1, 0]]),
 array([[1, 2, 1, 3],
       [1, 1, 2, 0]])]

Note that this happens to work also for longer numbers.

>>> n = np.array([462, 3, 15, 1, 0])
>>> pprint([np.reshape(nn, n.shape).astype(int) for nn in product(*map(str, n.ravel()))])
[array([4, 3, 1, 1, 0]),
 array([4, 3, 5, 1, 0]),
 array([6, 3, 1, 1, 0]),
 array([6, 3, 5, 1, 0]),
 array([2, 3, 1, 1, 0]),
 array([2, 3, 5, 1, 0])]

Approach 2: np.meshgrid

>>> import numpy as np
>>>
>>> n = np.array([[1, 12, 1, 3],
...               [1, 1, 12, 0]])
>>> 
>>> te = np.where(n>=10)
>>> dims = tuple(np.log10(n[te]).astype(int) + 1)
>>> 
>>> out = np.empty(dims + n.shape, dtype=n.dtype)
>>> out[...] = n
>>> out[(Ellipsis,) + te] = np.moveaxis(np.meshgrid(*(s//10**np.arange(i)[::-1]%10 for i, s in zip(dims, n[te])), indexing='ij'), 0, -1)
>>> 
>>> out
array([[[[1, 1, 1, 3],
         [1, 1, 1, 0]],

        [[1, 1, 1, 3],
         [1, 1, 2, 0]]],


       [[[1, 2, 1, 3],
         [1, 1, 1, 0]],

        [[1, 2, 1, 3],
         [1, 1, 2, 0]]]])

Upvotes: 1

Related Questions