Find formula for function of $n$ returning $0$ if $n$ is composite and $1$ if $n$ is prime Sample problem: 

Find an equation $\theta(n)$ for which $\theta(n)=\left\{ \begin{array} &0, \text{when } n\in \text{Composed} \\ n, \text{when } n\in \text{Prime} \end{array} \right.$

This problem is from the International Youth Math Challenge $2018$ and since they do not return marked sheets, I am unsure if my solution was correct. 
My final answer was: $$\theta (n)=n-n\cdot \text{sgn} \left(\prod_{i=1}^{\infty} |n-p_i|\right)$$ where $p_i$ is the $i^{\text{th}}$ prime number. This is all I could come up with and to be honest, I am not too happy with it, because I feel like I have basically chosen something that will only give the answer I want. Is this solution correct, mathematically? Is there a better solution? 
Note: $\text {sgn}(n)$ is the $\text{sign}$ or $\text{signum}$ function and $$\text{sgn}(n)=\left\{ \begin{array} &-1; \ \ n\lt 0\\ \ \ \ 0; \ \ n=0\\ \ \ \ 1; \ \ n\gt 0 \end{array} \right.$$
 A: I'm not sure what kind of functions are allowed but here is a similar one (might be equivalent after some small changes), the differences being it's finite and doesn't require ability to select primes:  
For any positive integer $p$, define this function
$$
f(n,p):= \left\lceil \frac{n-p\lfloor \frac{n}{p}\rfloor}{n} \right\rceil
$$
If $p$ divides $n$ then $f(n,p)=0$, otherwise $n-p\lfloor n/p\rfloor\neq 0$ so $f(n,p)=1$.  
You can then use this to make the following:
$$
\theta(n):= n - n\prod_{p=2}^{n-1}f(n,p)
$$
If $n$ is composite then one of the $p$'s will make the product $0$ and hence $\theta(n)=n$. Otherwise $n$ is prime and the product is $1$, giving $\theta(n)=0$.
A: Good try; but there's something that needs fixing. When $n$ is composite, the product diverges to $\infty$; so you should define $\text {sgn} (\infty) = 1$. 
A: Instead of the sign function, you could potentially use the Kronecker delta, which is defined as
$$\delta_{mn}=\begin{cases}
1 & \text{if }n=m\\
0 & \text{if }n\neq m
\end{cases}$$
It basically compares two numbers and gives $1$ if there is a match and $0$ otherwise. By summing over such Kronecker delta's, you could build:
$$\theta(n)=n\sum_{i=1}^{\infty}\delta_{np_i}$$
(where $p_i$ is the $i$-th prime number). However, I'm wondering if this would be accepted because it kind of bypasses the question of "checking if $n$ is prime". Both our formulas are just a nice rewording of the "text-form" formula given in the question, so I'm not certain this was the kind of answer that was expected.
A: Use wilsons theorem:
$$(n-1)!\equiv -1 (mod n)$$
This means that
$$\frac{(n-1)!+1}{n}$$ is an integer for prime n.
So we get that
$$\pi(\frac{(n-1)!+1}{n})$$
is an integer multiple of pi for prime n and a non-integer multiple of p for composite n. Using this we apply the cosine function to the above equation.
$$cos(\pi(\frac{(n-1)!+1}{n}))$$
Now for prime n the above is -1 or 1. For composite n the above is between -1 and 1. We want the above function to be postive, without allowing it to have a value greater than 1. So we can square it. We also want it to be 1 for prime n and 0 for composite n, so we can take the floor function. This gives us
$$\lfloor{cos^2(\pi(\frac{(n-1)!+1}{n}))}\rfloor$$
So we have a function that is zero for compostite n, and 1 for prime n, to get to what we want we simply multiply it by n. So
$$\theta(n)=\lfloor{cos^2(\pi(\frac{(n-1)!+1}{n}))}\rfloor$$
