changing the boolean values of an array according to a formula for the indices

Question

I want to create a 64 components array showing all the squares in which the two rooks of an empty chessboard could move from their current position. So far I am doing it with for and while loops.

I first create a function just to better visualize the board:

import numpy as np
def from_array_to_matrix(v):
    m=np.zeros((8,8)).astype('int')
    for row in range(8):
        for column in range(8): 
            m[row,column]=v[row*8+column]
    return m

and here I show how I actually build the array:

# positions of the two rooks

a=np.zeros(64).astype('int')
a[15] = 1
a[25] = 1

print from_array_to_matrix(a)

# attack_a will be all the squares where they could move in the empty board
attack_a=np.zeros(64).astype('int')
for piece in np.where(a)[0]:
    j=0
    square=piece+j*8
    while square<64:
        attack_a[square]=1
        j+=1
        square=piece+j*8
    j=0
    square=piece-j*8
    while square>=0:
        attack_a[square]=1
        j+=1
        square=piece-j*8
    j=0
    square=piece+j
    while square<8*(1+piece//8):
        attack_a[square]=1
        j+=1
        square=piece+j
    j=0
    square=piece-j
    while square>=8*(piece//8):
        attack_a[square]=1
        j+=1
        square=piece-j

print attack_a
print from_array_to_matrix(attack_a)

I have been advised to avoid for and while loops whenever it is possible to use other ways, because they tend to be time consuming. Is there any way to achieve the same result without iterating the process with for and while loops ? Perhaps using the fact that the indices to which I want to assign the value 1 can be determined by a function.

Answer 1

There are a couple of different ways to do this. The simplest thing is of course to work with matrices.

But you can vectorize operations on the raveled array as well. For example, say you had a rook at position 0 <= n < 64 in the linear array. To set the row to one, use integer division:

array[8 * (n // 8):8 * (n // 8 + 1)] = True

To set the column, use modulo:

array[n % 8::8] = True

You can convert to a matrix using reshape:

matrix = array.reshape(8, 8)

And back using ravel:

array = martix.ravel()

Or reshape:

array = matrix.reshape(-1)

Setting ones in a matrix is even simpler, given a specific row 0 <= m < 8 and column 0 <= n < 8:

matrix[m, :] = matrix[:, n] = True

Now the only question is how to vectorize multiple indices simultaneously. As it happens, you can use a fancy index in one axis. I.e, the expression above can be used with an m and n containing multiple elements:

m, n = np.nonzero(matrix)
matrix[m, :] = matrix[:, n] = True

You could even play games and do this with the array, also using fancy indexing:

n = np.nonzero(array)[0]
r = np.linspace(8 * (n // 8), 8 * (n // 8 + 1), 8, False).T.ravel()
c = np.linspace(n % 8, n  % 8 + 64, 8, False)
array[r] = array[c] = True

Using linspace allows you to generate multiple sequences of the same size simultaneously. Each sequence is a column, so we transpose before raveling, although this is not required.

Answer 2

Use reshaping to convert 1-D array to 8x8 2-D matrix and then numpy advance indexing to select rows and columns to set to 1:

import numpy as np
def from_array_to_matrix(v):
    return v.reshape(8,8)

# positions of the two rooks    
a=np.zeros(64).astype('int')
a[15] = 1
a[25] = 1

a = from_array_to_matrix(a)

# attack_a will be all the squares where they could move in the empty board
attack_a=np.zeros(64).astype('int')
attack_a = from_array_to_matrix(attack_a)

#these two lines replace your for and while loops
attack_a[np.where(a)[0],:] = 1
attack_a[:,np.where(a)[1]] = 1

output:

a:

[[0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 1]
 [0 0 0 0 0 0 0 0]
 [0 1 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0]]

attack_a:

[[0 1 0 0 0 0 0 1]
 [1 1 1 1 1 1 1 1]
 [0 1 0 0 0 0 0 1]
 [1 1 1 1 1 1 1 1]
 [0 1 0 0 0 0 0 1]
 [0 1 0 0 0 0 0 1]
 [0 1 0 0 0 0 0 1]
 [0 1 0 0 0 0 0 1]]