heyyyyy

mail@pastecode.io avatar
unknown
python
a year ago
1.1 kB
0
Indexable
Never
def remove_ones(arr):
    """
    Removes rows and columns of a 2D numpy array that only contain ones, unless there is a different value in the adjacent row or column.
    Inputs:
        arr: 2D numpy array
    Outputs:
        new_arr: 2D numpy array with specified rows and columns removed
    """
    rows_to_keep = []
    cols_to_keep = []
    
    for i in range(arr.shape[0]):
        if np.all(arr[i,:] == 1):
            if i > 0 and not np.all(arr[i-1,:] == 1):
                rows_to_keep.append(i)
            elif i < arr.shape[0]-1 and not np.all(arr[i+1,:] == 1):
                rows_to_keep.append(i)
        else:
            rows_to_keep.append(i)
            
    for j in range(arr.shape[1]):
        if np.all(arr[:,j] == 1):
            if j > 0 and not np.all(arr[:,j-1] == 1):
                cols_to_keep.append(j)
            elif j < arr.shape[1]-1 and not np.all(arr[:,j+1] == 1):
                cols_to_keep.append(j)
        else:
            cols_to_keep.append(j)
            
    new_arr = arr[np.ix_(rows_to_keep, cols_to_keep)]
    
    return new_arr