Moves to make all the elements in a array equal in Python

user_15753271 0
Viewed 386

I came across a problem in which the input is as follows:

5
1 1 1 1 6

And the expected output is 4

Basically what we are trying to do is print the minimum number of moves it will require to make all the 5 values equal to each other. One move means reducing a location and incrementing another location. If it is not possible to make them all equal, we print -1.

I tried the below approach:

def solution(N, W):
    counter=0
    W.sort()
    k=W[int(N/2)]
    for i in range(0,N):
        counter=counter+abs(W[i]-k)
        
    if N%2==0:
        tempC=0
        k=W[int(N/2)-1]
      
        for i in range(0,N):
            tempC=tempC+abs(A[i]-k)
        
        counter=min(counter,tempC)
    
    return counter

and am getting 5 as the answer. Kindly share what your function to achieve this would be.

math python
1 Answers

Lets see how logic works with your input.

5 1 1 1 1 6

1. If 1 2 1 3 3 this case possible then finally it show look like this 2 2 2 2 2. What are the things we get from this result Sum(INITAL_LIST) is equal to SUM(FINAL_LIST), this is 1st condition, if this hold pattern is possible.

2. Among all index-value some of them is going to leave some value some will take. Decrement of one-index and Increment of another-index is taken as one step, so we will consider only decrement case. Those who are leaving and finally become equal to 2. So total-step is equal to some of decremented index value.

Here I use vectorization properties of numpy for easy operation.

CODE :

import numpy as np

def solution(N, W):
    if sum(W)%N !=0:
        return -1
    
    eq = sum(W)//N
    step = np.array(W)-eq
    step_sum = np.sum(step[step>0])
    return step_sum


if __name__ == '__main__':
    var_, list_ = input().split(maxsplit=1)
    print(solution(int(var_), list(int(i) for i in list_.split())))

WITHOUT NUMPY :

Update :

step = np.array(W)-eq
step_sum = np.sum(step[step>0])

To :

step = [i-eq for i in W]
step_sum = sum(i for i in step if i>0)

OUTPUT :

5 1 1 1 1 6
4
user_148677300
Related