A chromatic version of Lagrange's four squares theorem (Q2514253)

Let \(\mathcal{S}\) be the sequence of squares of the natural numbers. For each \(K\geqslant 1\) let \[ \mathcal{S}=\bigcup_{1\leqslant i\leqslant K}\mathcal{S}_i \] be a partition of sequence \(\mathcal{S}\) into subsequences \(\{\mathcal{S}_1,\mathcal{S}_2,\dots,\mathcal{S}_K\}\). By \(s(K)\) denote the smallest integer such that for each sufficiently large integer \(n\) there exists some \(i\in\{1,2,\dots, K\}\) with the property that \(n\) can be expressed as a sum of no more than \(s(K)\) squares, all belonging to \(\mathcal{S}_i\). The authors prove that \(s(K)\leqslant c_\varepsilon K^{2+\varepsilon}\) for an arbitrary positive \(\varepsilon\) and an arbitrary \(K\geqslant 1\) with some positive quantity \(c_\varepsilon\) depending only on \(\varepsilon\).