Search code examples
pythonarraysnumpyvectorizationxor

Shifting all element positions of an array in Python given by XOR operator


I would like to know if there is a way to move all elements of a numpy array without iterating over each entry. The shift I desire is to relabel the indices by a fixed XOR operation, in the following form:

import numpy as np 

N = 2

x = np.arange(2**(2 * N)).reshape(2**N, 2**N)
z = np.zeros((2**N, 2**N))

k = 1

for i in range(2**N):
    for j in range(2**N):
        z[i][j] = x[i ^ k][j ^ k]

The problem I have is that latter I wish to take huge values of N, which becomes a bottleneck if we wish to iterate over each entry. Any advice on how to perform the move in a single shot will be very much appreciated.


Solution

  • There is simply no reason to use a loop here at all. You are not shifting anything, and the problem is completely separable across the axes:

    x = np.arange(2**(2 * N)).reshape(2**N, 2**N)
    z = x[np.arange(2**N)[:, None] ^ k, np.arange(2**N) ^ k]
    

    But what is x[a, b] for some a, b to begin with? From the original definition, you can see that it's a * 2**N + b. You can therefore plug in a = np.arange(2**N)[:, None] ^ k and b = np.arange(2**N) ^ k to avoid having to generate anything but the final result:

    idx = np.arange(2**N) ^ k
    z = (idx[:, None] << N) + idx
    

    It's generally faster to substitute << N for * 2**N.

    Allocation

    None of the solutions shown here pre-allocate z. However, if you were to do that, as in the original question, you should be careful about the type.

    z = np.zeros((2**N, 2**N))
    

    This creates an array of floats by default, which is likely not what you want.

    z = np.zeros((2**N, 2**N), dtype=int)
    

    Adding an explicit dtype make the array into integers. However, the simplest method available to you is likely

    z = np.zeros_like(x)
    

    Since you are planning on filling in every single element, you don't need to waste time filling it with zeros first. np.empty presents a better option in this case:

    z = np.empty((2**N, 2**N), dtype=int)
    

    OR

    z = np.empty_like(x)