Python - Fast way to split positive integer into array of ones

Viewed 123

I have an array of shape (100000, 1) with each element in the array is a positive integer and not greater than 6.

My goal is to convert each element into ones and place these ones in a new matrix of shape (100000, 6).

For example,

Input

X = np.array([[6],
              [2],
              [1],
              ...,
              [5],
              [4],
              [3]])
# shape of X is (100000, 1)

Output

Y = np.array([[1, 1, 1, 1, 1, 1],
              [1, 1, 0, 0, 0, 0],
              [1, 0, 0, 0, 0, 0],
              [       ...      ],
              [1, 1, 1, 1, 1, 0],
              [1, 1, 1, 1, 0, 0],
              [1, 1, 1, 0, 0, 0]])
# shape of Y is (100000, 6)

Is there any method that can achieve this without looping? Any help would be appreciated.

4 Answers

One way using numpy.flip with cumsum:

max_ = 6
np.flip(np.flip(np.eye(max_)[X.ravel()-1], 1).cumsum(1), 1)

Output:

array([[1., 1., 1., 1., 1., 1.],
       [1., 1., 0., 0., 0., 0.],
       [1., 0., 0., 0., 0., 0.],
       [1., 1., 1., 1., 1., 0.],
       [1., 1., 1., 1., 0., 0.],
       [1., 1., 1., 0., 0., 0.]])

Benchmark with 100k:

x_large = np.random.randint(1, 7, 100000)
max_ = 6
%timeit np.flip(np.flip(np.eye(max_)[x_large.ravel()-1], 1).cumsum(1), 1)
# 6.71 ms ± 68.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Proposing this solution as well. Chris' solution seems to perform slightly better on my computer. You can shave off a few ms by removing the (...) * 1 which simply converts booleans to integers.

>>> a = np.random.randint(1, 7 , 1000000).reshape(1000000, 1)
>>> %timeit (np.cumsum(np.tile(a, 6) >= 0, axis=1) <= a) * 1
75.9 ms ± 802 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

This seems to provide the same result and shave off a few other ms.

>>> %timeit (np.ones((1000000, 6)).cumsum(1) <= a).astype(int)
76 ms ± 695 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

If you do not convert to integer, then:

>>> %timeit (np.ones((1000000, 6)).cumsum(1) <= a)
73.8 ms ± 3.75 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

I've modified ApplePie's answer to get a faster method

Benchmark on 1M random samples

n_lines = 1000000
x_large = np.random.randint(1, 7, n_lines)
max_ = 6

# Chris
%timeit np.flip(np.flip(np.eye(max_)[x_large.ravel()-1], 1).cumsum(1), 1)
>>> 32.6 ms ± 706 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

x_large_shp = x_large.reshape((n_lines, 1))

# ApplePie
%timeit (np.cumsum(np.tile(x_large_shp, max_) >= 0, axis=1) <= x_large_shp) * 1
>>> 59.5 ms ± 615 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

# mine
%timeit np.repeat(np.arange(1,7),n_lines).reshape(6,-1).T <= x_large_shp * 1
>>> 22.9 ms ± 542 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

ApplePie builds a matrix of the following kind

[[1,2,3,4,5,6],
 [1,2,3,4,5,6],
 ...
 [1,2,3,4,5,6]]

And compare it, item by item with a broadcast version of the integer matrix

He builds the first matrix with the code snippet

np.cumsum(np.tile(a, 6) >= 0, axis=1)

Turns out cumsum is a bit slow, its faster to use repeat

np.repeat(np.arange(1,7),n_lines).reshape(6,-1).T

On my machine this cut more than half the time to run the original answer

You can use map method to do the job. In your case, it works much faster than a loop. Here's a basic use:

# One way to convert your array elements into array of 1 and 0
def myFunc(a):
    if a[0] == 1:
        return [1,0,0,0,0,0]
    elif a[0] == 2:
        return [1,1,0,0,0,0]
    elif a[0] == 3:
        return [1,1,1,0,0,0]
    elif a[0] == 4:
        return [1,1,1,1,0,0]
    elif a[0] == 5:
        return [1,1,1,1,1,0]
    else:
        return [1,1,1,1,1,1]


arr = [[6], [2],[3],[1]]

brr = list(map(myFunc, arr))
print(brr)
# prints: [[1, 1, 1, 1, 1, 1], [1, 1, 0, 0, 0, 0], [1, 1, 1, 0, 0, 0], [1, 0, 0, 0, 0, 0]]
Related