Weird Sum (Codeforces 1648A)
Tags: Math, matrices

Egor has a table of size $$$n \times m$$$, with lines numbered from $$$1$$$ to $$$n$$$ and columns numbered from $$$1$$$ to $$$m$$$. Each cell has a color that can be presented as an integer from $$$1$$$ to $$$10^5$$$.Let us denote the cell that lies in the intersection of the $$$r$$$-th row and the $$$c$$$-th column as $$$(r, c)$$$. We define the manhattan distance between two cells $$$(r_1, c_1)$$$ and $$$(r_2, c_2)$$$ as the length of a shortest path between them where each consecutive cells in the path must have a common side. The path can go through cells of any color. For example, in the table $$$3 \times 4$$$ the manhattan distance between $$$(1, 2)$$$ and $$$(3, 3)$$$ is $$$3$$$, one of the shortest paths is the following: $$$(1, 2) \to (2, 2) \to (2, 3) \to (3, 3)$$$. Egor decided to calculate the sum of manhattan distances between each pair of cells of the same color. Help him to calculate this sum.

Input: The first line contains two integers $$$n$$$ and $$$m$$$ ($$$1 \leq n \le m$$$, $$$n \cdot m \leq 100\,000$$$) — number of rows and columns in the table.Each of next $$$n$$$ lines describes a row of the table. The $$$i$$$-th line contains $$$m$$$ integers $$$c_{i1}, c_{i2}, \ldots, c_{im}$$$ ($$$1 \le c_{ij} \le 100\,000$$$) — colors of cells in the $$$i$$$-th row.

Output: Print one integer — the the sum of manhattan distances between each pair of cells of the same color.

Note: In the first sample there are three pairs of cells of same color: in cells $$$(1, 1)$$$ and $$$(2, 3)$$$, in cells $$$(1, 2)$$$ and $$$(2, 2)$$$, in cells $$$(1, 3)$$$ and $$$(2, 1)$$$. The manhattan distances between them are $$$3$$$, $$$1$$$ and $$$3$$$, the sum is $$$7$$$.

Sample Inputs:
2 3
1 2 3
3 2 1

3 4
1 1 2 2
2 1 1 2
2 2 1 1

4 4
1 1 2 3
2 1 1 2
3 1 2 1
1 1 2 1

Sample Outputs:
7

76

129

Solution (Python):
def main():
    (r, c) = map(int, input().split())
    m = [[] for _ in range(r)]
    for i in range(r):
        m[i] = list(map(int, input().split()))
    i_righe = dict()
    i_col = dict()
    for i in range(r):
        for j in range(c):
            elem = m[i][j]
            if elem in i_righe:
                i_righe[elem].append(i)
            else:
                i_righe[elem] = [i]
    for j in range(c):
        for i in range(r):
            elem = m[i][j]
            if elem in i_col:
                i_col[elem].append(j)
            else:
                i_col[elem] = [j]
    s = 0
    for ii in i_righe:
        length = len(i_righe[ii])
        if length > 1:
            for i in range(1, length):
                s += i_righe[ii][i] * i
                s -= i_righe[ii][i - 1] * (length - i)
    for ii in i_col:
        length = len(i_col[ii])
        if length > 1:
            for i in range(1, length):
                s += i_col[ii][i] * i
                s -= i_col[ii][i - 1] * (length - i)
    print(s)
main()

Solution (Python):
import sys
input = lambda : sys.stdin.readline()
from collections import defaultdict
(N, M) = map(int, input().split())
A = []
for _ in range(N):
    A.append(list(map(int, input().split())))
(R, C) = (defaultdict(list), defaultdict(list))
for i in range(N):
    for j in range(M):
        R[A[i][j]].append(i)
        C[A[i][j]].append(j)

def deal(D):
    ans = 0
    for (k, v) in D.items():
        n = len(v)
        if n < 2:
            continue
        v.sort()
        for (i, d) in enumerate(v):
            ans += d * (i - (n - i - 1))
    return ans
ans = deal(R)
ans += deal(C)
print(ans)

Solution (Python):
from collections import defaultdict

def main():
    (n, m) = list(map(int, input().split()))
    grid = []
    for _ in range(n):
        grid.append(list(map(int, input().split())))
    stat = defaultdict(list)
    for i in range(n):
        for j in range(m):
            stat[grid[i][j]].append((i, j))
    ans = 0
    for (k, v) in stat.items():
        v.sort(key=lambda x: x[0])
        s = 0
        for (i, x) in enumerate(v):
            ans += x[0] * i - s
            s += x[0]
        v.sort(key=lambda x: x[1])
        s = 0
        for (i, x) in enumerate(v):
            ans += x[1] * i - s
            s += x[1]
    print(ans)
if __name__ == '__main__':
    main()