Simulated annealing for magic square-of-squares puzzle
This is much more clear thanks to your edit.

Why do you store this as a list of squares? Can you store it as a unique list of positive integers, and square them later? That would make this massively simpler.

Another thing that would make this massively simpler is if you're able to guarantee that the list is maintained in some kind of order. In the example code below I do not assume this, so have to sort it myself. Maintaining order should be easy because the last_square mutation will not change order.


  Will my function always have the behaviour I've specified?


Sounds like you need to write some unit tests :) You will have to mock your random function to make them predictable.


  I care about checking if np.sqrt(c) != int(np.sqrt(c)) because a number whose square root is an integer is a square number.


Fine, but... you said that there was already a guarantee that these numbers are positive squares.

I tried to put together a vectorized version of this that does not loop. Rule number one when writing efficient Numpy code is to attempt to avoid loops. This solution roots, flattens and sorts the data, then looks for the closest gap. I do not know if it will perform better than yours, nor whether this is the most efficient implementation of this approach.

The performance will be influenced by the length of your data and the sparsity of your numbers. For highly sparse numbers the looping approach will probably do better. For highly dense numbers the approach below will probably do better.

def last_square_new(x):
    assert (x &gt; 0).all()
    x = x.copy()

    # The n-dimensional random index. Cannot yet be used as an index.
    nd_index = np.random.default_rng().integers(
        np.zeros(x.ndim), x.shape
    )
    # What will the index of the chosen element be once the array is flattened?
    flat_index = nd_index[0]*x.shape[-1] + nd_index[1]
    # Make this usable as an index.
    nd_index = nd_index[:, np.newaxis].tolist()

    xsq = (np.sqrt(x) + 0.5).astype(np.uint64)     # Sqrt with safe conversion to int
    chosen = xsq[nd_index][0]                      # The randomly chosen element's sqrt
    xsq = xsq.flatten()                            # Flatten the root array
    s_indices = xsq.argsort()                      # Indices that would sort the root array
    flat_index = s_indices[s_indices][flat_index]  # Move the flat index to its sorted position
    xsq = xsq[s_indices][:flat_index+1]            # Sort the root array and truncate

    d_indices = np.arange(xsq.size-1)[np.diff(xsq) &gt; 1]  # Which indices have gaps?
    if d_indices.size != 0:                              # Are there any gaps?
        gap_index = d_indices[-1]                        # Index of the closest gap (low side)
        best = xsq[gap_index+1] - 1                      # Best missing root
        x[nd_index] = best**2                            # Assign it to the output

    return x