The following code
import numpy as np
W = 5 # any odd positive integer
L = 7 # any positive integer
M = (W-1)//2
arr = np.random.choice(range(W),L) - M
generates an array of length L
with numbers ranging from -W
to W
. Now I want to amend this array such that its sum is 0 (while keeping the restriction that elements are within -W
to W
).
I want to achieve this by adding increments of +1 (if the initial sum is negative) or increments of -1 (if the initial sum is positive) to the elements of this array, starting from the first element of the list (and then the second, third, etc. should I need to). For example, if the initial sum is -2 (i.e. negative) and the array is [2,0,-1,...]
then I cannot increment by the first element by +1 (it's already at max W=2
), and thus I need to increment the second element two times, resulting in [2,2,-1,...]
.
Three samples:
[-1, 0, 0, -1, 2, 0, 2] ---> [-2, -1, 0, -1, 2, 0, 2]
or
[-2, -2, 1, -2, 0, -1, 2] ---> [ 2, -2, 1, -2, 0, -1, 2]
or
[ 2, 2, 2, 2, 2, 2, 2] ---> [-2, -2, -2, 0, 2, 2, 2]
It's easy to write a few for/while loops which do this, but is there a more efficient way (perhaps a numpy function)?
Aucun commentaire:
Enregistrer un commentaire