Triangles in a tetrahedron The goal of this challenge is to extend the OEIS sequence A334581.

Number of ways to choose \$3\$ points that form an equilateral triangle from the \$\binom{n+2}{3}\$ points in a regular tetrahedral grid of side length \$n\$.

The sequence begins 0, 0, 4, 24, 84, 224, 516, 1068, ...
Example
For example, if we take side lengths of \$n=4\$, as in the example below, there are \$\binom{4+2}{3} = 20\$ vertices and \$\binom{20}{3} = 1140\$ ways to choose three of them. Of these, there are \$84\$ ways to choose \$3\$ vertices that form an equilateral triangle:

Challenge
This fastest-code challenge will have you compute as many terms of the sequence [A334581(0), A334581(1), ...] that you can in ten minutes, as measured on my machine.
Rust, \$A(1), \dotsc, A(1375)\$ in 10 minutes

Unofficial score on Ryzen 7 1800X (8 cores/16 threads).  Build with cargo build --release and run with time target/release/tetrahedron n to compute \$A(1), \dotsc, A(n)\$.

This runs in \$O(n^4)\$ time.  (So to estimate a good value of \$n\$ for your CPU, first time it for some smaller \$n\$, then multiply that \$n\$ by a factor of \$\left(\frac{600\,\mathrm{s}}{t}\right)^{1/4}\$.)

How it works

Any triangle that fits inside a tetrahedron of minimal side \$k \le n\$ may be translated inside a tetrahedron of side \$n\$ in exactly \$\binom{n - k + 3}{3}\$ ways.  This means we only need to find it in one position, leaving six free parameters.  Two of these parameters may be computed from the other four (up to a sign choice) if the triangle is to be equilateral, so we only need to loop over an \$O(n^4)\$ space.

src/main.rs

use rayon::prelude::*;

fn get_counts(n: i64, a0: i64) -&gt; Vec&lt;i64&gt; {
    let mut c = vec![0; n as usize];
    let a0a0 = a0 * a0;
    for a1 in if a0 == 0 { 1 } else { -n + 1 }..n {
        let d = a0a0 + a1 * a1;
        let m = n - a0.abs() - a1.abs();
        for a2 in if m &gt; 0 {
            -n + 2 - (m &amp; 1)..n
        } else {
            -n - m + 2..n + m
        }
        .step_by(2)
        {
            let d = d + a2 * a2;
            let r = 2 * (a0a0 - d);
            if r == 0 {
                continue;
            }
            for b0 in a0..n {
                let pp = d * (3 * d - 4 * (a0a0 + b0 * (b0 - a0)));
                if pp &lt; 0 {
                    break;
                }
                let p = (pp as f64).sqrt() as i64;
                if p * p != pp {
                    continue;
                }
                let q = 2 * a0 * b0 - d;
                let mut check = |p: i64| {
                    let b1r = p * a2 + q * a1;
                    if b1r % r != 0 {
                        return;
                    }
                    let b1 = b1r / r;
                    let b2r = -p * a1 + q * a2;
                    if b2r % r != 0 {
                        return;
                    }
                    let b2 = b2r / r;
                    if (b0, b1, b2) &lt;= (a0, a1, a2) || b0 + b1 + b2 &amp; 1 != 0 {
                        return;
                    }
                    let t = 0.max(a0 + a1 + a2).max(b0 + b1 + b2)
                        + 0.max(-a0 - a1 + a2).max(-b0 - b1 + b2)
                        + 0.max(-a0 + a1 - a2).max(-b0 + b1 - b2)
                        + 0.max(a0 - a1 - a2).max(b0 - b1 - b2);
                    if t &gt;= 2 * n {
                        return;
                    }
                    c[t as usize / 2] += 1;
                };
                check(p);
                if p != 0 {
                    check(-p);
                }
            }
        }
    }
    c
}

fn add_vec(c0: Vec&lt;i64&gt;, c1: Vec&lt;i64&gt;) -&gt; Vec&lt;i64&gt; {
    c0.into_iter().zip(c1).map(|(x0, x1)| x0 + x1).collect()
}

fn main() {
    let n = std::env::args().skip(1).next().expect("missing argument");
    let n = n.parse().expect("not an integer");
    let counts = (0..n)
        .into_par_iter()
        .map(|a0| get_counts(n, a0))
        .reduce(|| vec![0; n as usize], add_vec);
    let (mut d0, mut d1, mut d2, mut d3) = (0, 0, 0, 0);
    for (i, x) in (1..).zip(counts) {
        d3 += x;
        d2 += d3;
        d1 += d2;
        d0 += d1;
        println!("{} {}", i, d0);
    }
}


Cargo.toml

[package]
name = "tetrahedron"
version = "0.1.0"
authors = ["Anders Kaseorg &lt;andersk@mit.edu&gt;"]
edition = "2018"

[dependencies]
rayon = "1.3.0"


Try it online!  (Parallelism removed for TIO.)