Skip to content

3426. Manhattan Distances of All Arrangements of Pieces

  • Time: $O(k \cdot \log\texttt{kMod})$
  • Space: $O(1)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Solution {
 public:
  int distanceSum(int m, int n, int k) {
    // For each distance d, where 1 < d < m, there are `m - d` ways to choose
    // the two columns that the two pieces are on. For each of the two pieces,
    // there are `n` ways to choose the row that the piece is on.
    // Therefore, the contribution of row differences is
    //   sum(d * (m - d) * n^2), where 1 < d <= m - 1
    // = n^2 * sum(d * m - d^2)
    // = n^2 * (d * m * (m - 1) / 2 - m * (m - 1) * (2m - 1) / 6)
    // = n^2 * (m^3 - m) / 6
    // Similarly, the contribution of column differences is
    //   m^2 * (n^3 - n) / 6
    const int rowContrib = 1L * n * n * (1L * m * m * m - m) / 6 % kMod;
    const int colContrib = 1L * m * m * (1L * n * n * n - n) / 6 % kMod;
    return (rowContrib + colContrib) * nCk(m * n - 2, k - 2) % kMod;
  }

 private:
  static constexpr int kMod = 1'000'000'007;

  long nCk(int n, int k) {
    long res = 1;
    for (int i = 1; i <= k; ++i)
      // By Fermat's little theorem.
      res = res * (n - i + 1) % kMod * modPow(i, kMod - 2) % kMod;
    return res;
  }

  long modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return x * modPow(x % kMod, (n - 1)) % kMod;
    return modPow(x * x % kMod, (n / 2)) % kMod;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution {
  public int distanceSum(int m, int n, int k) {
    final long rowContrib = 1L * n * n * (1L * m * m * m - m) / 6 % MOD;
    final long colContrib = 1L * m * m * (1L * n * n * n - n) / 6 % MOD;
    return (int) ((rowContrib + colContrib) % MOD * nCk(m * n - 2, k - 2) % MOD);
  }

  private static final int MOD = 1_000_000_007;

  private long nCk(int n, int k) {
    long res = 1;
    for (int i = 1; i <= k; ++i)
      // By Fermat's little theorem.
      res = res * (n - i + 1) % MOD * modPow(i, MOD - 2) % MOD;
    return res;
  }

  private long modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return x * modPow(x, n - 1) % MOD;
    return modPow(x * x % MOD, n / 2);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
class Solution:
  def distanceSum(self, m: int, n: int, k: int) -> int:
    # For each distance d, where 1 < d < m, there are `m - d` ways to choose
    # the two columns that the two pieces are on. For each of the two pieces,
    # there are `n` ways to choose the row that the piece is on.
    # Therefore, the contribution of row differences is
    #   sum(d * (m - d) * n^2), where 1 < d <= m - 1
    # = n^2 * sum(d * m - d^2)
    # = n^2 * (d * m * (m - 1) / 2 - m * (m - 1) * (2m - 1) / 6)
    # = n^2 * (m^3 - m) / 6
    # Similarly, the contribution of column differences is
    #   m^2 * (n^3 - n) / 6
    MOD = 1_000_000_007
    return (n**2 * (m**3 - m) // 6 +
            m**2 * (n**3 - n) // 6) * math.comb(m * n - 2, k - 2) % MOD