IMO Team Selection Test

We prove the following equivalent statement: we show that if $f : \mathbb{N} \to \mathbb{Z}$ is a function such that $(f(m) - f(n))(m - n)$ is always the square of an integer, then $f$ must be of the form $A^2x + B$ for integers $A, B$. First, since $p(f(n + p) - f(n))$ is a square for any prime $p$ and positive integer $n$, $p$ must divide $f(n + p) - f(n)$. (Similarly, we may prove $p^k \mid f(n + p^k) - f(n)$ by induction, but we will not need to do so.)

Lemma. If $p \mid f(a) - f(b)$ but $p \nmid a - b$ for some $a, b$, then $p \mid f(m) - f(n)$ for all $m, n$. (In other words, $f$ is either injective or constant modulo $p$.)
Proof. For each integer $r$, define $S_r = \{x : f(x) \equiv r \pmod{p}\}$, which must be a union of residue classes modulo $p$ (restricted to the positive integers). Then we need to show that either (1) for every $r$, $S_r$ covers at most one residue class (and thus exactly one), or (2) for some $r$, $S_r$ covers all residue classes.

Suppose otherwise; then there exists $r$ such that $S = S_r$ covers between 2 and $p-1$ residues (so $p \ge 3$). If $n \notin S$ (i.e. $f(n) \not\equiv r \pmod p$), then for any $s \in S$,
$$(n-s)(f(n)-f(s)) \equiv (n-s)(f(n)-r) \not\equiv 0 \pmod{p}$$
is a nonzero quadratic residue. Hence $\left(\frac{n-s}{p}\right) = \left(\frac{f(n)-r}{p}\right) \neq 0$.
Now let $T = \{t : \left(\frac{f(t)-r}{p}\right) = 1\}$ and $U = \{u : \left(\frac{f(u)-r}{p}\right) = -1\}$, so $T, U$ partition $\mathbb{N} \setminus S$. By the previous paragraph, $\left(\frac{a}{p}\right) = 1$ for any $a = t-s$ in the difference set $T-S$; similarly, $\left(\frac{b}{p}\right) = -1$ for any $b \in U-S$. This immediately upper bounds $|T-S|$ and $|U-S|$ by $\frac{p-1}{2}$, the number of (nonzero) quadratic residues and nonresidues.
But $|T| + |U| = p - |S| \ge 1$, so if $|T| \ge |U|$ (so $T$ is nonempty), then Cauchy-Davenport yields
$$\frac{p-1}{2} \ge |T-S| \ge \min(p, |T| + |S| - 1) \ge \min\left(p, \frac{p-|S|}{2} + |S| - 1\right) = \min\left(p, \frac{p + (|S| - 2)}{2}\right),$$
contradicting $|S| \ge 2$. The case $|U| \ge |T|$ is analogous. $\square$

Note that $f(n+1) - f(n)$ is always a square. If $f$ is nonconstant, $\gcd_{n \ge 1}(f(n+1) - f(n)) \ne 0$ must be a square itself (say $g^2$, with $g > 0$). If $g=1$, then $f$ is nonconstant (and thus injective, by the lemma) modulo every prime $p$. In particular, $p \nmid f(n+1) - f(n)$ for all $n$ and $p$, which forces $f(x+1) - f(x) \equiv 1 \implies f(x) \equiv x+d$ for some constant $d$. Otherwise, if $g > 1$, $f'(x) \equiv \frac{f(x)-f(1)}{g^2}$ has $(f'(m) - f'(n))(m-n)$ always a square and $\gcd_{n \ge 1}(f'(n+1) - f'(n)) = 1$, so $f(x) \equiv g^2(x+d) + c$ for some constants $d, c$.

It follows that $f(x) \equiv A^2x + B$ ($A, B \in \mathbb{Z}$) are the only possible solutions, which indeed all work (note that we get $f$ constant when $A=0$).

We give an alternate proof of the lemma. Define $r$ and $S = S_r$ as in the previous proof.
Since $p \ge 3$, there exists a smallest quadratic nonresidue $\alpha \in [2, p-1]$ modulo $p$. In particular, $\alpha-1 \in [1, p-2]$ is a quadratic residue. Now fix two distinct residues $x, y$ (mod $p$) in $S$. We claim that for $N = x+\alpha(y-x)$, we have $f(N) \equiv r \pmod p$. Suppose otherwise; then $f(N) \not\equiv r \equiv f(x) \equiv f(y)$. But then $[f(N)-f(x)][N-x] \equiv [f(N)-r]\alpha(y-x)$ and $[f(N)-f(y)][N-y] \equiv [f(N)-r](\alpha-1)(y-x)$ must both be nonzero squares, forcing $\left(\frac{\alpha}{p}\right) = \left(\frac{\alpha-1}{p}\right)$, which is absurd.
Hence $x+\alpha(y-x) = (1-\alpha)x + \alpha y$ lies in $S$ (interpret relations modulo $p$ where clearly appropriate) for any distinct residues $x, y \in S$; of course, it also does when $x \equiv y \in S$. This condition is affine, so for convenience, fix distinct residues $a, b \in S$ and define $A = (S-a)(b-a)^{-1}$. Then $0, 1 \in A$ and we still have $(1-\alpha)x + \alpha y \in A$ whenever $x, y \in A$.
By plugging in $x \equiv 0$, we get $\alpha A \subseteq A$, and from $y \equiv 0$, we get $(1-\alpha)A \subseteq A$. Therefore (noting that $p \nmid \alpha, 1-\alpha$)
$$A + A \equiv (1-\alpha)(1-\alpha)^{p-2}A + \alpha\alpha^{p-2}A \subseteq (1-\alpha)A + \alpha A \subseteq A,$$
so $A+1 \subseteq A \implies A = \mathbb{Z}/p\mathbb{Z} \implies S = \mathbb{Z}/p\mathbb{Z}$, as desired.

We present yet another proof of the lemma. The previous proof shows that if $x, y \in S$, then $x + \alpha(y-x) \in S$. We claim that $N = x + (\alpha+1)(y-x)$ lies in $S$ as well. Indeed, the only way $[f(N)-f(y)][N-y] \equiv [f(N)-r]\alpha(y-x)$ and $[f(N)-f(x+\alpha(y-x))][N-(x+\alpha(y-x))] \equiv [f(N)-r](y-x)$ can both be quadratic residues is if $p \mid f(N)-r$, so $N \in S$.
But $[x+(\alpha+1)(y-x)] - [x+\alpha(y-x)] = y-x$, so we conclude that $x+k\alpha(y-x)$ and $x+k(\alpha+1)(y-x)$ lie in $S$ for all $k \ge 0$. Since $p \nmid \alpha$, $x+k\alpha(y-x)$ covers all residues modulo $p$, so $S = \mathbb{Z}/p\mathbb{Z}$, as desired.