Kazakhstan FAIO 2024 · Correction by IOAI Training

F. Gravitational Matrix

This page gives the full approach. Read it once you have searched.

Metric: exact matrix on hidden tests
Evaluated on: self-generated tests (official tests are hidden)
Our score: 200 random checks vs numpy

We solved the problem ourselves: downloaded the data, ran the experiments and measured the score below. The statement and data belong to their organisers.

The problem in brief

Read two matrices, $A$ of size $a\times b$ and $B$ of size $b\times c$ , with all dimensions up to 300 and real entries, and print their product $AB$ , an $a\times c$ matrix. That is the whole task: a direct exercise on the operation at the heart of every neural network, where each layer computes $W x + b$ for a whole batch at once. The only real difficulty is speed, with a 1 second time limit.

How you are scored

Hidden tests compare your matrix with the exact product (the tolerance is not stated; we print floats with full precision via repr). The README gives no sample, so we test against numpy on 200 random shapes and on a hand-computed $2\times2$ case.

The key idea

Definition. Entry $(i, j)$ of the product is the dot product of row $i$ of $A$ with column $j$ of $B$ :

(AB)_{ij} = \sum_{k=1}^{b} A_{ik} B_{kj}.

The inner dimensions must agree ( $A$ has $b$ columns, $B$ has $b$ rows), and the result takes the outer ones, $a\times c$ . For example

\begin{pmatrix}1&2\\3&4\end{pmatrix}\begin{pmatrix}5&6\\7&8\end{pmatrix} = \begin{pmatrix}1\cdot5+2\cdot7 & 1\cdot6+2\cdot8\\ 3\cdot5+4\cdot7 & 3\cdot6+4\cdot8\end{pmatrix} = \begin{pmatrix}19&22\\43&50\end{pmatrix}.

A second reading of the same formula. Row $i$ of $AB$ is a linear combination of the rows of $B$ , with weights taken from row $i$ of $A$ :

(AB)_{i,:} = \sum_{k} A_{ik}\, B_{k,:}.

This is the view behind "a layer mixes its inputs": every output row is a weighted blend of the input rows.

Cost. Each of the $ac$ entries needs $b$ multiply-adds, so the product costs $abc$ operations: up to $300^3 = 2.7\cdot10^7$ . A modern CPU does that in a few milliseconds, but a Python interpreter executes each s += A[i][k] * B[k][j] as several bytecode instructions, roughly 40 nanoseconds each round. That puts a pure Python triple loop right at the time limit.

Running time of three implementations of the same product, log-log scale

On our machine, the triple loop takes 1.08 s for $n = 300$ , and the row-combination version 0.96 s, both on the edge of the limit before even counting the time to read and print 90 000 numbers. numpy's @, which calls an optimised BLAS routine (compiled, vectorised, cache-friendly), takes 7 milliseconds, 150 times faster. All three curves have slope 3 on the log-log plot: same $O(n^3)$ algorithm, very different constant factors.

The full solution

import sys, numpy as np
t = sys.stdin.read().split()
a, b, c = map(int, t[:3])
v = np.array(t[3:3 + a*b + b*c], dtype=float)
A, B = v[:a*b].reshape(a, b), v[a*b:].reshape(b, c)
for row in A @ B:
    print(" ".join(repr(float(x)) for x in row))

Reading every token at once with split() does not depend on line breaks, and it is much faster than calling input() 600 times. solution.py also contains the two pure Python versions, used for testing and as a fallback if numpy were unavailable on the judge.

Results

check	result
200 random shapes up to $12\times12\times12$ : triple loop, row combination, numpy	all agree (`np.allclose`)
hand-computed $2\times2$ product	`19 22 / 43 50`
$n=300$ : triple loop / row combination / numpy	1.08 s / 0.96 s / 0.007 s

What to remember

$(AB)_{ij}$ is row $i$ of $A$ dotted with column $j$ of $B$ ; equivalently, each row of $AB$ is a combination of the rows of $B$ . Inner dimensions must match.
The cost is $abc$ multiply-adds: doubling every dimension multiplies the time by 8.
Vectorise: in Python, push the inner loops into numpy (BLAS). The algorithm is the same, the constant factor is 100 times smaller. This is why deep learning frameworks phrase everything as matrix products.
Read input in bulk (sys.stdin.read().split()) when it is large.

Full code (solution.py)

"""FAIO 2024 qualification, problem F: Gravitational Matrix -- matrix product A (a x b) times B (b x c).

stdin : a b c, then a rows of b reals, then b rows of c reals
stdout: the a x c product, one row per line

C[i][j] = sum_k A[i][k] B[k][j]: a*b*c multiply-adds (27 million for 300 x 300 x 300). A pure
Python triple loop takes several seconds, over the 1 s limit; numpy's BLAS does it in milliseconds.
A fallback without numpy (row combination with zip) is included.

Run: python solution.py < input.txt ; python solution.py --check ; python solution.py --figure
"""
import sys


def read(text):
    t = text.split()
    a, b, c = int(t[0]), int(t[1]), int(t[2])
    vals = list(map(float, t[3:3 + a * b + b * c]))
    A = [vals[i * b:(i + 1) * b] for i in range(a)]
    B = [vals[a * b + k * c: a * b + (k + 1) * c] for k in range(b)]
    return A, B


def matmul_naive(A, B):
    """Textbook definition: three nested loops (i, j, k)."""
    a, b, c = len(A), len(B), len(B[0])
    C = [[0.0] * c for _ in range(a)]
    for i in range(a):
        for j in range(c):
            s = 0.0
            for k in range(b):
                s += A[i][k] * B[k][j]
            C[i][j] = s
    return C


def matmul_rows(A, B):
    """Row i of C is a linear combination of the rows of B with weights A[i][k] (loop order i, k, j)."""
    c = len(B[0])
    C = []
    for row in A:
        acc = [0.0] * c
        for w, brow in zip(row, B):
            if w:
                acc = [x + w * y for x, y in zip(acc, brow)]
        C.append(acc)
    return C


def matmul(A, B):
    try:
        import numpy as np
        return (np.array(A) @ np.array(B)).tolist()
    except ImportError:
        return matmul_rows(A, B)


def fmt(v):
    s = repr(float(v))
    return s[:-2] if s.endswith(".0") else s


def solve():
    A, B = read(sys.stdin.read())
    out = sys.stdout
    out.write("\n".join(" ".join(fmt(v) for v in row) for row in matmul(A, B)) + "\n")


def check():
    import random, time
    import numpy as np
    rng = random.Random(0)
    for _ in range(200):
        a, b, c = (rng.randint(1, 12) for _ in range(3))
        A = [[rng.uniform(-100, 100) for _ in range(b)] for _ in range(a)]
        B = [[rng.uniform(-100, 100) for _ in range(c)] for _ in range(b)]
        ref = np.array(A) @ np.array(B)
        for f in (matmul_naive, matmul_rows, matmul):
            assert np.allclose(f(A, B), ref)
    assert matmul([[1, 2], [3, 4]], [[5, 6], [7, 8]]) == [[19, 22], [43, 50]]
    print("200 random shapes: naive, row-combination and numpy agree; [[1,2],[3,4]][[5,6],[7,8]] = [[19,22],[43,50]]")
    res = []
    for n in (50, 100, 200, 300):
        A = [[rng.uniform(-1, 1) for _ in range(n)] for _ in range(n)]
        row = [n]
        for f in (matmul_naive, matmul_rows, matmul):
            t = time.perf_counter(); f(A, A); row.append(time.perf_counter() - t)
        res.append(row)
        print(f"n={n}: naive {row[1]:.3f}s, row-combination {row[2]:.3f}s, numpy {row[3]:.4f}s")
    return res


def figure():
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    res = check()
    n = [r[0] for r in res]
    fig, ax = plt.subplots(figsize=(7.5, 4.5))
    for j, (nm, col) in enumerate([("triple loop (i, j, k)", "#d95f02"), ("row combination with zip", "#7570b3"),
                                   ("numpy @ (BLAS)", "#1b9e77")], start=1):
        ax.loglog(n, [r[j] for r in res], "o-", color=col, label=nm)
    ax.axhline(1, color="crimson", ls="--", lw=1.2); ax.text(52, 1.15, "time limit 1 s", color="crimson")
    ax.set_xlabel("n (multiplying two n x n matrices)"); ax.set_ylabel("seconds (this machine)")
    ax.set_title("Same n³ multiply-adds, very different speeds"); ax.legend()
    fig.tight_layout(); fig.savefig("figures/timing.png", dpi=150, facecolor="white")


if __name__ == "__main__":
    if "--check" in sys.argv:
        check()
    elif "--figure" in sys.argv:
        figure()
    else:
        solve()