CODEWITHSUNDEEP
pythonarraysnumpy
Jack Wetherell
Jack Wetherell

Reputation: 585

Numpy: Conserving sum in average over two arrays of integers

I have two arrays of positive integers A and B that each sum to 10:

I want to write a code (that will work for a general size of the array and the sum) to calculate the array C who is also a array of positive integers that also sums to 10 that is the closest to the element-wise average as possible:

Any value can be adjusted to make the sum correct, as long as all elements remain positive integers.

What should C = SOME CODE be?

Minimum reproducible example:

A = np.array([1,4,5])
B = np.array([5,5,0])
C = np.ceil((A + B) / 2).astype(int)
print(np.sum(C))

11

This should give 10.

Upvotes: 1

Views: 55

Answers (2)

Ehsan
Ehsan

Reputation: 12407

You can ceil/floor every other non-int element. This works for any shape/size and any sum value (in fact you do not need to know the sum at all. It is enough if A and B have same sum):

C = (A + B) / 2
C_c = np.ceil(C)
C_c[np.flatnonzero([C!=C.astype(int)])[::2]] -= 1

print(C_c.sum())
#10.0
print(C_c.astype(int))
#[3 4 3]

Upvotes: 2

Chrispresso
Chrispresso

Reputation: 4081

Ok so based off what you're saying, this could work:

C = ((a + b) / 2)  # array([3, 4, 2])
curr_sum = sum(C)  # 9
adjust_amount = sum(a) - curr_sum  # 10-9 = 1
if adjust_amount > 0:
    C[-1] += adjust_amount  # array([3, 4, 3])
# Otherwise if it's negative just grab the largest and subtract to ensure you still remain >0
else:
    C[np.argmax(C)] += adjust_amount

Upvotes: 1

Related Questions