通过调整负条目来限制 numpy 数组中的条目总和

Question

I have a numpy array containing positive and negative values, and I want to adjust the negative entries so that the sum is not negative, starting with the most negative entry. The maximum adjustment is to make a negative entry zero. I have an implementation using a loop, is there a way to do it using numpy array methods? Here is my code:

initial_values = np.asarray([50,-200,-180,110])
sorted_index = np.argsort(initial_values)

final_values = initial_values
for i, entry in enumerate(final_values[sorted_index]):
    ss = final_values.sum()
    if ss >= 0:
        break
    adjustment = max(entry, ss)
    final_values[sorted_index[i]] -= adjustment

print final_values

The starting array is [50,-200,-180,110], the answer in this case is [50, 0, -160, 110], so the most negative entry is set to zero, and then the next most negative entry is adjusted to make the sum zero.

Does anyone have a simpler, faster numpy based solution?

Answer 1

Here's one vectorized approach -

# Get a copy of input as the output
out = initial_values.copy()

# Get sorted indices
sorted_index = np.argsort(out)

# Mask of elements that would be made zero for sure and zero them
mask = out.sum() < out[sorted_index].cumsum()
out[sorted_index[mask]] = 0

# There might be one element left to make the sum absolutely zero. 
# Make it less negative to make the absolute sum zero.
out[sorted_index[np.where(mask)[0][-1]+1]] -= out.sum()

Sample run -

Function definitions -

In [155]: def vectorized(initial_values):
     ...:   out = initial_values.copy()
     ...:   sorted_index = np.argsort(out)
     ...:   mask = out.sum() < out[sorted_index].cumsum()
     ...:   out[sorted_index[mask]] = 0
     ...:   out[sorted_index[np.where(mask)[0][-1]+1]] -= out.sum()
     ...:   return out
     ...: 
     ...: def org_app(initial_values):  
     ...:   final_values = initial_values.copy()
     ...:   sorted_index = np.argsort(initial_values)
     ...:   for i, entry in enumerate(final_values[sorted_index]):
     ...:       ss = final_values.sum()
     ...:       if ss >= 0:
     ...:           break
     ...:       adjustment = max(entry, ss)
     ...:       final_values[sorted_index[i]] -= adjustment
     ...:   return final_values
     ...: 

Case #1 :

In [156]: initial_values
Out[156]: array([  50, -200, -180,  110])

In [157]: vectorized(initial_values)
Out[157]: array([  50,    0, -160,  110])

In [158]: org_app(initial_values)
Out[158]: array([  50,    0, -160,  110])

Case #2 :

In [163]: initial_values
Out[163]: array([ 50, -20, -14, -22, -15,   6, -21, -19, -17,   4,   5, -56])

In [164]: vectorized(initial_values)
Out[164]: array([ 50,   0, -14,   0, -15,   6,   0, -19, -17,   4,   5,   0])

In [165]: org_app(initial_values)
Out[165]: array([ 50,   0, -14,   0, -15,   6,   0, -19, -17,   4,   5,   0])

Runtime tests -

In [177]: initial_values = np.random.randint(-100,20,(50000))

In [178]: np.array_equal(vectorized(initial_values),org_app(initial_values))
Out[178]: True

In [179]: %timeit org_app(initial_values)
1 loops, best of 3: 2.08 s per loop

In [180]: %timeit vectorized(initial_values)
100 loops, best of 3: 5.7 ms per loop

---

Here's a slightly improved (lesser code and better runtime) version of the earlier proposed approach -

# Get a copy of input as the output
out = initial_values.copy()

# Get sorted indices
sorted_index = np.argsort(out)

# Last index in sorted indexed indices for setting elements in input array to 0's
idx = np.where(out.sum() < out[sorted_index].cumsum())[0][-1]

# Set until idx indexed into sorted_index in turn indexed into input array t0 0's
out[sorted_index[:idx+1]] = 0

# There might be one element left to make the sum absolutely zero. 
# Make it less negative to make the absolute sum zero.
out[sorted_index[idx+1]] -= out.sum()

Runtime tests -

In [18]: initial_values = np.random.randint(-100,20,(50000))

In [19]: %timeit vectorized(initial_values)
100 loops, best of 3: 5.58 ms per loop

In [20]: %timeit vectorized_v2(initial_values) # improved version
100 loops, best of 3: 5.4 ms per loop