Manhattan Distances of All Arrangements of Pieces
Source: leetcode
Difficulty: Hard
Tags: math, combinatorics

You are given three integers `m`, `n`, and `k`.


There is a rectangular grid of size `m  n` containing `k` identical pieces. Return the sum of Manhattan distances between every pair of pieces over all **valid arrangements** of pieces.


A **valid arrangement** is a placement of all `k` pieces on the grid with **at most** one piece per cell.


Since the answer may be very large, return it **modulo** `109 + 7`.


The Manhattan Distance between two cells `(xi, yi)` and `(xj, yj)` is `|xi - xj| + |yi - yj|`.


 


**Example 1:**



**Input:** m \= 2, n \= 2, k \= 2


**Output:** 8


**Explanation:**


The valid arrangements of pieces on the board are:


![](https://assets.leetcode.com/uploads/2024/12/25/4040example1.drawio)![](https://assets.leetcode.com/uploads/2024/12/25/untitled-diagramdrawio.png)


- In the first 4 arrangements, the Manhattan distance between the two pieces is 1\.
- In the last 2 arrangements, the Manhattan distance between the two pieces is 2\.


Thus, the total Manhattan distance across all valid arrangements is `1 + 1 + 1 + 1 + 2 + 2 = 8`.



**Example 2:**



**Input:** m \= 1, n \= 4, k \= 3


**Output:** 20


**Explanation:**


The valid arrangements of pieces on the board are:


![](https://assets.leetcode.com/uploads/2024/12/25/4040example2drawio.png)


- The first and last arrangements have a total Manhattan distance of `1 + 1 + 2 = 4`.
- The middle two arrangements have a total Manhattan distance of `1 + 2 + 3 = 6`.


The total Manhattan distance between all pairs of pieces across all arrangements is `4 + 6 + 6 + 4 = 20`.



 


**Constraints:**


- `1 <= m, n <= 105`
- `2 <= m * n <= 105`
- `2 <= k <= m * n`




C++ solution:
class Solution {
    long long mod_pow(long long base, int exp, int mod) {
        long long result = 1;
        base %= mod;
        while (exp > 0) {
            if (exp % 2 == 1) {
                result = (result * base) % mod;
            }
            base = (base * base) % mod;
            exp /= 2;
        }
        return result;
    }
public:
    int distanceSum(int m, int n, int k) {
        const int MOD = 1e9 + 7;
        long long inv6 = mod_pow(6, MOD-2, MOD);
        
        long long sum_x_part = (1LL * n * n % MOD) * m % MOD;
        sum_x_part = sum_x_part * ((1LL * m * m % MOD - 1 + MOD) % MOD) % MOD;
        
        long long sum_y_part = (1LL * m * m % MOD) * n % MOD;
        sum_y_part = sum_y_part * ((1LL * n * n % MOD - 1 + MOD) % MOD) % MOD;
        
        long long sum_total = (sum_x_part + sum_y_part) % MOD;
        sum_total = sum_total * inv6 % MOD;
        
        int a = m * n - 2;
        int b = k - 2;
        if (b < 0 || b > a) {
            return 0;
        }
        
        vector<long long> fact(a+1), inv_fact(a+1);
        fact[0] = 1;
        for (int i = 1; i <= a; ++i) {
            fact[i] = fact[i-1] * i % MOD;
        }
        inv_fact[a] = mod_pow(fact[a], MOD-2, MOD);
        for (int i = a-1; i >= 0; --i) {
            inv_fact[i] = inv_fact[i+1] * (i+1) % MOD;
        }
        
        long long comb = fact[a] * inv_fact[b] % MOD;
        comb = comb * inv_fact[a - b] % MOD;
        
        long long ans = sum_total * comb % MOD;
        return (int)ans;
    }
};

Go solution:
func distanceSum(m int, n int, k int) int {
    const MOD = 1000000007
    inv6 := 166666668 // Modular inverse of 6 modulo 1e9+7

    T := m * n
    if k < 2 || k > T {
        return 0
    }

    a := T - 2
    b := k - 2
    if a < 0 || b < 0 || a < b {
        return 0
    }

    term1 := (n % MOD) * (n % MOD) % MOD
    term1 = term1 * ((m - 1) % MOD) % MOD
    term1 = term1 * (m % MOD) % MOD
    term1 = term1 * ((m + 1) % MOD) % MOD

    term2 := (m % MOD) * (m % MOD) % MOD
    term2 = term2 * ((n - 1) % MOD) % MOD
    term2 = term2 * (n % MOD) % MOD
    term2 = term2 * ((n + 1) % MOD) % MOD

    sum_S := (term1 + term2) % MOD
    sum_S = sum_S * inv6 % MOD

    if a == 0 {
        return sum_S
    }

    fact := make([]int, a+1)
    fact[0] = 1
    for i := 1; i <= a; i++ {
        fact[i] = fact[i-1] * i % MOD
    }

    inv_fact := make([]int, a+1)
    inv_fact[a] = pow(fact[a], MOD-2, MOD)
    for i := a - 1; i >= 0; i-- {
        inv_fact[i] = inv_fact[i+1] * (i + 1) % MOD
    }

    comb := fact[a] * inv_fact[b] % MOD
    comb = comb * inv_fact[a-b] % MOD

    total := sum_S * comb % MOD
    return total
}

func pow(x, n, mod int) int {
    res := 1
    x = x % mod
    for n > 0 {
        if n%2 == 1 {
            res = (res * x) % mod
        }
        x = (x * x) % mod
        n /= 2
    }
    return res
}

Java solution:
class Solution {
    static final int MOD = 1000000007;
    static final int MAX = 100000;
    static long[] fact = new long[MAX + 1];
    static long[] invFact = new long[MAX + 1];
    static final int inv3 = 333333336;
    static final int inv2 = 500000004;

    static {
        fact[0] = 1;
        for (int i = 1; i <= MAX; i++) {
            fact[i] = fact[i - 1] * i % MOD;
        }
        invFact[MAX] = pow(fact[MAX], MOD - 2);
        for (int i = MAX - 1; i >= 0; i--) {
            invFact[i] = invFact[i + 1] * (i + 1) % MOD;
        }
    }

    private static long pow(long a, int b) {
        long res = 1;
        while (b > 0) {
            if ((b & 1) == 1) {
                res = res * a % MOD;
            }
            a = a * a % MOD;
            b >>= 1;
        }
        return res;
    }

    private static long comb(int a, int b) {
        if (a < 0 || b < 0 || a < b) {
            return 0;
        }
        return fact[a] * invFact[b] % MOD * invFact[a - b] % MOD;
    }

    public int distanceSum(int m, int n, int k) {
        int mn = m * n;
        if (k < 2 || k > mn) {
            return 0;
        }

        long mCube = (long) m * m % MOD * m % MOD;
        long mPart = (mCube - m + MOD) % MOD;
        long mTerm = (long) n * n % MOD;
        mTerm = mTerm * mPart % MOD;

        long nCube = (long) n * n % MOD * n % MOD;
        long nPart = (nCube - n + MOD) % MOD;
        long nTerm = (long) m * m % MOD;
        nTerm = nTerm * nPart % MOD;

        long sumManhattan = (mTerm + nTerm) % MOD;
        sumManhattan = sumManhattan * inv3 % MOD;

        long sumUnordered = sumManhattan * inv2 % MOD;

        long c = comb(mn - 2, k - 2);

        long ans = sumUnordered * c % MOD;

        return (int) ans;
    }
}

JavaScript solution:
var distanceSum = function(m, n, k) {
    const MOD = BigInt(1e9 +7);
    const inv3 = BigInt(333333336);
    const inv2 = BigInt(500000004);

    function powMod(base, exp, mod) {
        if (mod === 1n) return 0n;
        let result = 1n;
        base = base % mod;
        while (exp > 0n) {
            if (exp % 2n === 1n) {
                result = (result * base) % mod;
            }
            exp = exp / 2n;
            base = (base * base) % mod;
        }
        return result;
    }

    function computeComb(a, b, MOD) {
        if (a < 0 || b < 0 || a < b) return 0n;
        let fact = new Array(a+1);
        fact[0] = 1n;
        for (let i=1; i<=a; i++) {
            fact[i] = (fact[i-1] * BigInt(i)) % MOD;
        }
        let inv_fact = new Array(a+1);
        inv_fact[a] = powMod(fact[a], MOD - 2n, MOD);
        for (let i=a-1; i>=0; i--) {
            inv_fact[i] = (inv_fact[i+1] * BigInt(i+1)) % MOD;
        }
        return (fact[a] * inv_fact[b] % MOD) * inv_fact[a - b] % MOD;
    }

    let total = BigInt(m) * BigInt(n);
    let a = total - 2n;
    let b = BigInt(k) - 2n;
    if (a < b) return 0;

    let comb = computeComb(Number(a), Number(b), MOD);
    let sum_x = 0n;
    if (m >= 2) {
        let m_bi = BigInt(m), n_bi = BigInt(n);
        let term = (n_bi * n_bi) % MOD;
        term = (term * m_bi) % MOD;
        term = term * (m_bi * m_bi - 1n) % MOD;
        term = term * inv3 % MOD;
        sum_x = term * inv2 % MOD;
    }

    let sum_y = 0n;
    if (n >= 2) {
        let m_bi = BigInt(m), n_bi = BigInt(n);
        let term = (m_bi * m_bi) % MOD;
        term = (term * n_bi) % MOD;
        term = term * (n_bi * n_bi - 1n) % MOD;
        term = term * inv3 % MOD;
        sum_y = term * inv2 % MOD;
    }

    let total_sum = (sum_x + sum_y) % MOD;
    total_sum = (total_sum * comb) % MOD;
    return Number(total_sum);
};

Python solution:
class Solution:
    MOD = 10**9 + 7
    max_fact = 10**5
    fact = [1] * (max_fact + 1)
    for i in range(1, max_fact + 1):
        fact[i] = fact[i-1] * i % MOD
    inv_fact = [1] * (max_fact + 1)
    inv_fact[max_fact] = pow(fact[max_fact], MOD-2, MOD)
    for i in range(max_fact -1, -1, -1):
        inv_fact[i] = inv_fact[i+1] * (i+1) % MOD
    inv6 = pow(6, MOD-2, MOD)

    def distanceSum(self, m: int, n: int, k: int) -> int:
        a = m * n
        if k < 2 or k > a:
            return 0
        a_minus_2 = a - 2
        b = k - 2
        if a_minus_2 < b:
            return 0
        
        comb = self.fact[a_minus_2] * self.inv_fact[b] % self.MOD
        comb = comb * self.inv_fact[a_minus_2 - b] % self.MOD
        
        T = m % self.MOD
        T = T * (n % self.MOD) % self.MOD
        sum_mn = (m + n) % self.MOD
        T = T * sum_mn % self.MOD
        a_minus_1 = (a - 1) % self.MOD
        T = T * a_minus_1 % self.MOD
        
        res = comb * T % self.MOD
        res = res * self.inv6 % self.MOD
        return res

Ruby solution:
MOD = 10**9 + 7
INV6 = 166666668 # Modular inverse of 6 under MOD

def distance_sum(m, n, k)
  cells = m * n
  return 0 if k < 2 || k > cells

  a = cells - 2
  b = k - 2

  comb = if a < b
           0
         else
           max_fact = a
           fact = Array.new(max_fact + 1, 1)
           (2..max_fact).each { |i| fact[i] = fact[i-1] * i % MOD }

           inv_fact = Array.new(max_fact + 1, 1)
           inv_fact[max_fact] = fact[max_fact].pow(MOD - 2, MOD)
           (max_fact - 1).downto(0) { |i| inv_fact[i] = inv_fact[i+1] * (i+1) % MOD }

           fact[a] * inv_fact[b] % MOD * inv_fact[a - b] % MOD
         end

  sum_x = n % MOD
  sum_x = sum_x * n % MOD
  sum_x = sum_x * ( (m - 1) % MOD ) % MOD
  sum_x = sum_x * ( m % MOD ) % MOD
  sum_x = sum_x * ( (m + 1) % MOD ) % MOD
  sum_x = sum_x * INV6 % MOD

  sum_y = m % MOD
  sum_y = sum_y * m % MOD
  sum_y = sum_y * ( (n - 1) % MOD ) % MOD
  sum_y = sum_y * ( n % MOD ) % MOD
  sum_y = sum_y * ( (n + 1) % MOD ) % MOD
  sum_y = sum_y * INV6 % MOD

  total = (sum_x + sum_y) % MOD
  total * comb % MOD
end