import Mathlib
/-!
# International Mathematical Olympiad 2016, Problem 2
Find all positive integers $n$ for which each cell of an $n \times n$ table can
be filled with one of the letters $I$, $M$ and $O$ in such a way that:
* in each row and each column, one third of the entries are $I$, one third are
$M$ and one third are $O$; and
* in any diagonal, if the number of entries on the diagonal is a multiple of
three, then one third of the entries are $I$, one third are $M$ and one third
are $O$.
Note. The rows and columns of an $n \times n$ table are each labelled $1$ to $n$
in a natural order. Thus each cell corresponds to a pair of positive integers
$(i, j)$ with $1 \le i, j \le n$. For $n > 1$, the table has $4n - 2$ diagonals
of two types. A diagonal of the first type consists of all cells $(i, j)$ for
which $i + j$ is a constant, and a diagonal of the second type consists of all
cells $(i, j)$ for which $i - j$ is constant.
-/
namespace Imo2016P2
/-- The three letters `I`, `M`, `O`. -/
abbrev Letter := Fin 3
variable {n : ℕ}
/-- A finite set of cells is *balanced* under the coloring `f` if each of the
three letters occurs on exactly one third of its cells. -/
def Balanced (f : Fin n → Fin n → Letter) (cells : Finset (Fin n × Fin n)) : Prop :=
∀ c : Letter, (cells.filter (fun p ↦ f p.1 p.2 = c)).card = cells.card / 3
/-- The cells of row `r`. -/
def row (r : Fin n) : Finset (Fin n × Fin n) := Finset.univ.filter (fun p ↦ p.1 = r)
/-- The cells of column `k`. -/
def col (k : Fin n) : Finset (Fin n × Fin n) := Finset.univ.filter (fun p ↦ p.2 = k)
/-- A diagonal of the first type: the cells `(i, j)` with `i + j = s`. -/
def diagSum (s : ℕ) : Finset (Fin n × Fin n) :=
Finset.univ.filter (fun p ↦ (p.1 : ℕ) + (p.2 : ℕ) = s)
/-- A diagonal of the second type: the cells `(i, j)` with `i - j = d`. -/
def diagDiff (d : ℤ) : Finset (Fin n × Fin n) :=
Finset.univ.filter (fun p ↦ ((p.1 : ℕ) : ℤ) - ((p.2 : ℕ) : ℤ) = d)
/-- The coloring `f` of the `n × n` table satisfies all the required conditions:
every row and column is balanced, and every diagonal whose length is a multiple
of three is balanced. -/
def MODTable (f : Fin n → Fin n → Letter) : Prop :=
(∀ r : Fin n, Balanced f (row r)) ∧
(∀ k : Fin n, Balanced f (col k)) ∧
(∀ s : ℕ, 3 ∣ (diagSum (n := n) s).card → Balanced f (diagSum s)) ∧
(∀ d : ℤ, 3 ∣ (diagDiff (n := n) d).card → Balanced f (diagDiff d))
/- determine -/ abbrev solution : Set ℕ := sorry
theorem imo2016_p2 (n : ℕ) (hn : 0 < n) :
n ∈ solution ↔ ∃ f : Fin n → Fin n → Letter, MODTable f := sorry
end Imo2016P2
This problem does not yet have a complete formalized solution.