Manipulating rows and columns in an N,X,Y dimensional array in Python

Question

I'm trying to set N random rows and columns equal to a chosen value (in this case 1) in a multidimensional array of zeros of shape (N,X,Y) in Python.

Is there a way to do this more efficiently with indexing? So far I can only wrap my head around using a loop to set the value of chosen rows and columns.

Current approach:

# Import packages

import numpy as np
import matplotlib.pyplot as plt
import torch 

def create_batch(N):
    # Create an N by 64 by 64 array
    batch = np.zeros((N,64,64)) 
    # Generate the rows and columns you want to manipulate in each slice of N
    randoms = np.round(np.random.uniform(0,63,(N,2,5))).astype(int)
    # Set the chosen rows and columns to have value == 1
    for i in range(batch.shape[0]):
        batch[i,randoms[i,0],:] = 1; batch[i,:,randoms[i,1]] = 1;
    return batch

b = create_batch(3) # Create an array of shape (3,64,64)

# Visualise the column and row manipulation has worked
for i in range(b.shape[0]):
    plt.figure(); plt.imshow(b[i])

question from:https://stackoverflow.com/questions/65940741/manipulating-rows-and-columns-in-an-n-x-y-dimensional-array-in-python

Answer 1

Waitting for answers