Reshaping array with specified indices

Viewed 47

Is there any better way to do this? Like replacing that list comprehension with numpy functions? I'd assume that for a small number of elements, the difference is insignificant, but for larger chunks of data it takes too much time.

>>> rows = 3
>>> cols = 3
>>> target = [0, 4, 7, 8] # each value represent target index of 2-d array converted to 1-d
>>> x = [1 if i in target else 0 for i in range(rows * cols)]
>>> arr = np.reshape(x, (rows, cols))
>>> arr 
[[1 0 0]
 [0 1 0]
 [0 1 1]]
2 Answers

Another way:

shape = (rows, cols)
arr = np.zeros(shape)
arr[np.unravel_index(target, shape)] = 1

Since x comes from a range, you can index an array of zeros to set the ones:

x = np.zeros(rows * cols, dtype=bool)
x[target] = True
x = x.reshape(rows, cols)

Alternatively, you can create the proper shape up front and assign to the raveled array:

x = np.zeros((rows, cols), dtype=bool)
x.ravel()[target] = True

If you want actual zeros and ones, use a dtype like np.uint8 or whatever else suits your needs other than bool.

The approach shown here would apply even to your list example to make it more efficient. Even if you turned target into a set, you are performing O(N) lookups, with N = rows * cols. Instead, you only need M assignments with no lookups, with M = len(target):

x = [0] * (rows * cols)
for i in target:
    x[i] = 1
Related