Cost Equilibrium (Codeforces 1513E)
Tags: Math, combinatorics

An array is called beautiful if all the elements in the array are equal.You can transform an array using the following steps any number of times:   Choose two indices $$$i$$$ and $$$j$$$ ($$$1 \leq i,j \leq n$$$), and an integer $$$x$$$ ($$$1 \leq x \leq a_i$$$). Let $$$i$$$ be the source index and $$$j$$$ be the sink index.  Decrease the $$$i$$$-th element by $$$x$$$, and increase the $$$j$$$-th element by $$$x$$$. The resulting values at $$$i$$$-th and $$$j$$$-th index are $$$a_i-x$$$ and $$$a_j+x$$$ respectively.  The cost of this operation is $$$x \cdot |j-i| $$$.  Now the $$$i$$$-th index can no longer be the sink and the $$$j$$$-th index can no longer be the source.  The total cost of a transformation is the sum of all the costs in step $$$3$$$.For example, array $$$[0, 2, 3, 3]$$$ can be transformed into a beautiful array $$$[2, 2, 2, 2]$$$ with total cost $$$1 \cdot |1-3| + 1 \cdot |1-4| = 5$$$.An array is called balanced, if it can be transformed into a beautiful array, and the cost of such transformation is uniquely defined. In other words, the minimum cost of transformation into a beautiful array equals the maximum cost.You are given an array $$$a_1, a_2, \ldots, a_n$$$ of length $$$n$$$, consisting of non-negative integers. Your task is to find the number of balanced arrays which are permutations of the given array. Two arrays are considered different, if elements at some position differ. Since the answer can be large, output it modulo $$$10^9 + 7$$$.

Input: The first line contains a single integer $$$n$$$ ($$$1 \leq n \leq 10^5$$$) — the size of the array. The second line contains $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$ ($$$0 \le a_i \le 10^9$$$).

Output: Output a single integer — the number of balanced permutations modulo $$$10^9+7$$$.

Note: In the first example, $$$[1, 2, 3]$$$ is a valid permutation as we can consider the index with value $$$3$$$ as the source and index with value $$$1$$$ as the sink. Thus, after conversion we get a beautiful array $$$[2, 2, 2]$$$, and the total cost would be $$$2$$$. We can show that this is the only transformation of this array that leads to a beautiful array. Similarly, we can check for other permutations too.In the second example, $$$[0, 0, 4, 4]$$$ and $$$[4, 4, 0, 0]$$$ are balanced permutations.In the third example, all permutations are balanced.

Sample Inputs:
3
1 2 3

4
0 4 0 4

5
0 11 12 13 14

Sample Outputs:
6

2

120

Solution (Python):
import sys
input = sys.stdin.readline

def solve():
    n = int(input())
    a = list(map(int, input().split()))
    a.sort()
    s = sum(a)
    if s % n != 0:
        print(0)
        return
    F = [0] * (n + 1)
    FI = [0] * (n + 1)
    F[0] = 1
    MOD = int(1000000000.0 + 7)
    for i in range(1, n + 1):
        F[i] = F[i - 1] * i % MOD
    FI[n] = pow(F[n], MOD - 2, MOD)
    for i in range(n - 1, -1, -1):
        FI[i] = FI[i + 1] * (i + 1) % MOD
    z = s // n
    b = []
    c = []
    e = 0
    i = 0
    while i < n:
        v = a[i]
        j = i + 1
        while j < n and a[j] == v:
            j += 1
        if v < z:
            b.append(j - i)
        elif v > z:
            c.append(j - i)
        else:
            e += j - i
        i = j
    if len(b) == 1 and b[0] == 1 or (len(c) == 1 and c[0] == 1):
        r = F[n]
        for v in b:
            r = r * FI[v] % MOD
        for v in c:
            r = r * FI[v] % MOD
        r = r * FI[e] % MOD
        print(r)
    elif len(b) == 0:
        print(1)
    else:
        r = 2
        s1 = 0
        for v in b:
            r = r * FI[v] % MOD
            s1 += v
        r = r * F[s1] % MOD
        s2 = 0
        for v in c:
            r = r * FI[v] % MOD
            s2 += v
        r = r * F[s2] % MOD
        r = r * F[n] * FI[n - e] * FI[e] % MOD
        print(r)
solve()

Solution (Python):
import io, os
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline
M = 10 ** 9 + 7
permute = [1, 1]
for _ in range(100010):
    permute.append(len(permute) * permute[-1] % M)
MAXNUM = 100010
factor = [1] * (MAXNUM + 2)
factor[-1] = 0
storecomb = {}
for i in range(1, MAXNUM + 1):
    factor[i] = factor[i - 1] * i % M

def fastfrac(a, b):
    numb = pow(b, M - 2, M)
    return a % M * (numb % M) % M

def comb(n, k):
    if n < k:
        return 0
    if n == k:
        return 1
    if (n, k) in storecomb:
        return storecomb[n, k]
    num1 = factor[n]
    num2 = factor[k]
    num3 = factor[n - k]
    num = num2 * num3 % M
    output = fastfrac(num1, num)
    storecomb[n, k] = output
    return output

def main(t):
    n = int(input())
    arr = list(map(int, input().split()))
    tot = sum(arr)
    fre = {}
    for num in arr:
        if num not in fre:
            fre[num] = 0
        fre[num] += 1
        if fre[num] == n:
            print(1)
            return
    if tot % n > 0:
        print(0)
        return
    small = 0
    equal = 0
    large = 0
    for i in range(n):
        if arr[i] > tot // n:
            large += 1
        elif arr[i] == tot // n:
            equal += 1
        else:
            small += 1
    if min(small, large) == 1:
        ans = permute[n]
    else:
        ans = comb(n, equal) * permute[small] % M
        ans = ans * permute[large] % M
        ans = ans * permute[equal] % M
        ans = ans * 2 % M
    for key in fre:
        if fre[key] > 1:
            ans = fastfrac(ans, permute[fre[key]])
    print(ans)
T = 1
t = 1
while t <= T:
    main(t)
    t += 1

Solution (Python):
from sys import stdin, stdout, maxsize
R = lambda : stdin.readline().strip()
RL = lambda f=None: list(map(f, R().split(' '))) if f else list(R().split(' '))
output = lambda x: stdout.write(str(x) + '\n')
output_list = lambda x: output(' '.join(map(str, x)))
M = int(1000000000.0) + 7

def fact(n, M):
    ans = 1
    for i in range(1, n + 1):
        ans = ans * i % M % M
    return ans

def mod_inv(n, M):
    return pow(n, M - 2, M)

def npr(n, r, M):
    return fact(n, M) * mod_inv(fact(n - r, M), M) % M

def ncr(n, r, M):
    return npr(n, r, M) * mod_inv(fact(r, M), M) % M

def permute(d, M):
    n = sum(list(d.values()))
    ans = fact(n, M)
    deno = 1
    for i in d.values():
        deno = deno * fact(i, M) % M
    ans = ans * mod_inv(deno, M) % M
    return ans

def get_deno(d, M):
    ans = 1
    deno = 1
    for i in d.values():
        deno = deno * fact(i, M) % M
    ans = ans * mod_inv(deno, M) % M
    return ans
n = int(R())
a = RL(int)
if n == 1 or max(a) == min(a):
    print(1)
    exit()
if sum(a) % n:
    print(0)
    exit()
mean = int(sum(a) / n)
small = dict()
large = dict()
cntmean = 0
for i in a:
    if i < mean:
        if i in small:
            small[i] += 1
        else:
            small[i] = 1
    elif i > mean:
        if i in large:
            large[i] += 1
        else:
            large[i] = 1
    else:
        cntmean += 1
ans = ncr(n, cntmean, M)
if sum(list(small.values())) == 1 or sum(list(large.values())) == 1:
    ans = fact(n, M)
    deno = mod_inv(fact(cntmean, M), M) * get_deno(small, M) % M
    deno = deno * get_deno(large, M) % M
    ans = ans * deno % M
    print(ans)
    exit()
ans = ans * (permute(small, M) * permute(large, M) % M) % M
ans = 2 * ans % M
print(ans)