Manhattan distances of all arrangements of pieces

Math Combinatorics

35.2% acceptance

Feb 25, 2026

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|.

Solution

Rust

Time O(2^n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn distance_sum(m: i32, n: i32, k: i32) -> i32 {
    const MOD: i64 = 1_000_000_007;
    let m = m as i64;
    let n = n as i64;
    let k = k as i64;
    let total = m * n;

    fn pow_mod(mut b: i64, mut e: i64, md: i64) -> i64 {
      let mut r = 1i64;
      b %= md;
      while e > 0 {
        if e & 1 == 1 { r = r * b % md; }
        b = b * b % md;
        e >>= 1;
      }
      r
    }

    // Precompute factorials and inverse factorials
    let max_n = (total + 1) as usize;
    let mut fact = vec![1i64; max_n + 1];
    for i in 1..=max_n { fact[i] = fact[i-1] * i as i64 % MOD; }
    let mut inv_fact = vec![1i64; max_n + 1];
    inv_fact[max_n] = pow_mod(fact[max_n], MOD - 2, MOD);
    for i in (0..max_n).rev() { inv_fact[i] = inv_fact[i+1] * (i+1) as i64 % MOD; }

    let comb = |a: i64, b: i64| -> i64 {
      if b < 0 || b > a { return 0; }
      fact[a as usize] * inv_fact[b as usize] % MOD * inv_fact[(a-b) as usize] % MOD
    };

    // Sum of |xi - xj| over all pairs from m rows, weighted by C(total-2, k-2) for rows
    // contribution from rows: sum over pairs (r1,r2), |r1-r2| * C(total-2, k-2) * n^2... hmm
    // Actually: distance = |xi - xj| + |yi - yj|
    // By linearity, total = (sum_row_distances) * C(n*n combs...) + ...
    // Better: separate row and col contributions
    // For row contribution: sum over all arrangements of k pieces of sum_{p<q} |row_p - row_q|
    // = C(total, k) * ... no, need to think carefully.
    //
    // The sum over ALL valid arrangements of k pieces of sum_{all pairs (i,j)} Manhattan dist
    // = sum over all valid arrangements of [sum_pairs |ri-rj| + sum_pairs |ci-cj|]
    // = (row contribution) + (col contribution)
    //
    // Row contribution = sum over all pairs of cells (c1,c2) of |row(c1)-row(c2)| *
    //   (number of arrangements containing both c1 and c2)
    //   = sum_{c1 != c2} |row(c1)-row(c2)| * C(total-2, k-2)
    //
    // Col contribution similarly.
    //
    // sum_{c1 != c2} |row(c1)-row(c2)| = n^2 * sum_{r1,r2=0}^{m-1} |r1-r2| * (pairs of columns per row pair)
    // Wait: c1 and c2 are cells (r1,col1) and (r2,col2). |row(c1)-row(c2)| = |r1-r2|.
    // Sum over all ordered pairs of distinct cells = sum_{r1,r2,c1,c2, (r1,c1)!=(r2,c2)} |r1-r2|
    // = n^2 * sum_{r1!=r2} |r1-r2| * (1 per col pair... actually n^2 col pairs for each row pair)
    //   + 0 (same row: |r1-r2|=0)
    // = n^2 * sum_{r1!=r2} |r1-r2|
    //
    // sum_{r1,r2=0}^{m-1} |r1-r2| (ordered, including r1=r2 which contributes 0):
    // = 2 * sum_{0<=r1<r2<=m-1} (r2-r1) = 2 * m*(m-1)*(m+1)/6... let me compute.
    // sum_{d=1}^{m-1} d * (m-d) * 2 = 2 * sum d(m-d)

    // sum_{r1,r2=0}^{m-1} |r1-r2| = 2 * sum_{d=1}^{m-1} d*(m-d) = m*(m-1)*(m+1)/3
    // (closed-form formula, O(1) per dimension)
    let row_sum = {
      let mm = m % MOD;
      let mm1 = (m - 1) % MOD;
      let mp1 = (m + 1) % MOD;
      mm * mm1 % MOD * mp1 % MOD * pow_mod(3, MOD - 2, MOD) % MOD
    };
    let col_sum = {
      let nm = n % MOD;
      let nm1 = (n - 1) % MOD;
      let np1 = (n + 1) % MOD;
      nm * nm1 % MOD * np1 % MOD * pow_mod(3, MOD - 2, MOD) % MOD
    };

    // Total = (n^2 * row_sum + m^2 * col_sum) * C(total-2, k-2) / 2
    // Wait: we're summing over UNORDERED pairs (c1,c2). 
    // Ordered pairs give twice the sum.
    // So: total_dist = C(total-2, k-2) * (n^2 * row_sum + m^2 * col_sum) / 2
    // But row_sum includes r1=r2 (which gives 0) and is ordered.
    // n^2 * row_sum counts ordered pairs of cells (c1,c2) with any c1,c2 (including c1=c2).
    // Since |row(c1)-row(c2)| = 0 when c1=c2, those don't contribute.
    // Ordered pairs contribute twice the unordered sum.
    // Actually:
    // sum over unordered pairs {c1,c2} of |row(c1)-row(c2)| 
    // = (1/2) * sum over ordered pairs (c1,c2), c1!=c2 of |row(c1)-row(c2)|
    // = (1/2) * n^2 * sum_{r1!=r2} |r1-r2| (since for each row pair, n^2 column pairs)
    // Hmm but row_sum already includes r1=r2 contributing 0, so:
    // sum_{ordered (r1,r2)} |r1-r2| = row_sum. And for r1=r2: contributes 0.
    // So ordered cell pair sum = n^2 * row_sum.
    // Unordered cell pair sum of |row| = n^2 * row_sum / 2.
    // Similarly for cols.
    // Total = C(total-2, k-2) * (n^2 * row_sum / 2 + m^2 * col_sum / 2)
    //       = C(total-2, k-2) * (n^2 * row_sum + m^2 * col_sum) / 2

    let inv2 = pow_mod(2, MOD - 2, MOD);
    let c_k2 = comb(total - 2, k - 2);
    let result = (n % MOD * n % MOD * row_sum % MOD + m % MOD * m % MOD * col_sum % MOD) % MOD
      * c_k2 % MOD * inv2 % MOD;
    result as i32
  }
}