Make Equal (Codeforces 1931B)
Tags: Greedy, greedy

There are $$$n$$$ containers of water lined up, numbered from left to right from $$$1$$$ to $$$n$$$. Each container can hold any amount of water; initially, the $$$i$$$-th container contains $$$a_i$$$ units of water. The sum of $$$a_i$$$ is divisible by $$$n$$$.You can apply the following operation any (possibly zero) number of times: pour any amount of water from the $$$i$$$-th container to the $$$j$$$-th container, where $$$i$$$ must be less than $$$j$$$ (i.e. $$$i<j$$$). Any index can be chosen as $$$i$$$ or $$$j$$$ any number of times.Determine whether it is possible to make the amount of water in all containers the same using this operation.

Input: The first line of the input contains a single integer $$$t$$$ ($$$1 \le t \le 10^4$$$) — the number of test cases. Then the descriptions of the test cases follow.The first line of each test case contains a single integer $$$n$$$ ($$$1 \le n \le 2 \cdot 10^5$$$) — the number of containers with water.The second line of each test case contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ ($$$0 \le a_i \le 10^9$$$) — the amounts of water in the containers. It is guaranteed that the sum of $$$a_i$$$ in each test case does not exceed $$$2 \cdot 10^9$$$. Also, the sum of $$$a_i$$$ is divisible by $$$n$$$.It is guaranteed that the sum of $$$n$$$ over all test cases in the input does not exceed $$$2 \cdot 10^5$$$.

Output: Output $$$t$$$ lines, each of which is the answer to the corresponding test case. As the answer, output "YES" if it is possible to make the amount of water in all containers the same using the described operation. Otherwise, output "NO".You can output each letter in any case (lowercase or uppercase). For example, the strings "yEs", "yes", "Yes", and "YES" will be accepted as a positive answer.

Note: In the third test case of the example ($$$a=[4, 5, 2, 1, 3]$$$), you can proceed as follows:  pour $$$1$$$ unit of water from the first vessel to the fourth, then $$$a=[3, 5, 2, 2, 3]$$$;  pour $$$1$$$ unit of water from the second vessel to the third, then $$$a=[3, 4, 3, 2, 3]$$$;  pour $$$1$$$ unit of water from the second vessel to the fourth, then $$$a=[3, 3, 3, 3, 3]$$$. 

Sample Inputs:
6
1
43
2
1 3
5
4 5 2 1 3
3
1 2 3
7
4 5 5 0 6 4 4
7
6 5 5 1 3 4 4

Sample Outputs:
YES
NO
YES
NO
NO
YES

Solution (Python):
import sys
from os import path

def input():
    return sys.stdin.readline().strip()

def solve(length, array):
    if len(array) < 2:
        return print('YES')
    total_num = sum(array)
    if total_num % length != 0:
        return print('NO')
    average_value = int(total_num / length)
    avg_value_in_hand = 0
    for current_value in array:
        c_eval = current_value - average_value
        avg_value_in_hand += c_eval
        if avg_value_in_hand < 0:
            return print('NO')
    return print('YES')

def main():
    testcases = int(input())
    for _ in range(testcases):
        n = int(input())
        array = list(map(int, input().split()))
        solve(n, array)
if __name__ == '__main__':
    if path.exists('input.txt'):
        sys.stdin = open('input.txt')
        sys.stdout = open('output.txt', 'w')
    main()

Solution (Python):
def make_equal(test_cases):
    answers = []
    for (n, a) in test_cases:
        target = sum(a) // n
        possible = 'YES'
        excess_sum = 0
        for water in a:
            if water > target:
                excess_sum += water - target
            elif water < target:
                excess_sum -= target - water
                if excess_sum < 0:
                    possible = 'NO'
                    break
        answers.append(possible)
    return answers
t = int(input())
test_cases = []
for _ in range(t):
    n = int(input())
    a = list(map(int, input().split()))
    test_cases.append((n, a))
answers = make_equal(test_cases)
for answer in answers:
    print(answer)

Solution (Python):
import sys

def get_ints():
    return map(int, sys.stdin.readline().strip().split())

def get_list():
    return list(map(int, sys.stdin.readline().strip().split()))

def get_list_string():
    return list(map(str, sys.stdin.readline().strip().split()))

def get_string():
    return sys.stdin.readline().strip()

def get_int():
    return int(sys.stdin.readline().strip())

def get_print_int(x):
    sys.stdout.write(str(x) + '\n')

def get_print(x):
    sys.stdout.write(x + '\n')

def get_print_int_same(x):
    sys.stdout.write(str(x) + ' ')

def get_print_same(x):
    sys.stdout.write(x + ' ')

def solve():
    for _ in range(get_int()):
        n = get_int()
        arr = get_list()
        opt = sum(arr) // n
        water = 0
        ans = 'YES'
        for i in range(n):
            if arr[i] > opt:
                water += arr[i] - opt
            elif arr[i] < opt:
                diff = opt - arr[i]
                if water < diff:
                    ans = 'NO'
                    break
                water -= diff
        get_print(ans)
solve()