I came with this proof and I found some other proofs online, but mine is different and I want to see if I made any mistakes.
Problem:
Suppose $|G| = p^3$, where $p$ is a prime. Show that either $|Z(G)|=p$ or $G$ is abelian.
Case 1: $G$ is not abelian
We have the class equation
$$|G| = |Z(G)| + \sum_{g\in G\setminus Z(G) } \frac{|G|}{|\mathrm{Cent}(g)|}$$
where $Z(G)$ is the center of the group, and $|\mathrm{Cent} (g)|$ is the centralizer of $g$.
if we solve for $|Z(G)|$ we get
$$|Z(G)| = |G| - \sum_{g\in G\setminus Z(G) } \frac{|G|}{|\mathrm{Cent}(g)|} = p^3 \bigg( 1-\sum \frac{1}{\mathrm{Cent}(g)} \bigg)$$
Also, since $Z(G) \leq G$, by Lagrange's theorem, $|Z(G)|$ divides $|G| = p^3 = ppp$
Therefore, we have 4 possibilities: $|Z(G)| = \{1,p,p^2,p^3\}$
It can't be $p^3$ since that implies that $G$ is abelian, contradicting the initial assumption. There's another theorem that states that if $|G|$ is a power of a prime number, then the center of $G$ contains nonidentity elements, so it can't be $1$.
This leaves us with $p$ or $p^2$.
If the size is $p^2$, we can go back to the class equation and obtain $$p^3 = p^2 + \sum\frac{p^3}{|\mathrm{Cent}(g)|}$$ $$p\bigg( 1 - \sum\frac{p}{\mathrm{Cent}(g)}\bigg) = 1$$ $$\sum\frac{p}{\mathrm{Cent}(g)} = \frac{p-1}{p}$$
This number has to be an integer greater than or equal to $0$, and the only option is then $p=1$, but we already discarded that option, and $|Z(G)| \neq p^2$
Therefore, $|Z(G)| = p$ if $G$ is nonabelian.
If it is abelian then $G \setminus Z(G) = \emptyset$ so the sum is over no elements giving $0$, and $Z(G)=G \leq G$.
Suppose that $G$ is non-abelian.
We know that $Z(G)\leq G$.
By Lagrange's Theorem $|Z(G)|$ must divide $|G|$.
Since $|G|=p^{3}$ the only possibilities are $1, p, p^{2}, p^{3}$.
$|Z(G)|\neq p^{3}$ because otherwise we will have $Z(G)=G$ but $G$ is non-abelian.
$|Z(G)|\neq p^{2}$ also because otherwise we will have the order of the factor group by the center as $|G/Z(G)|=|G|/|Z(G)|= p^{3}/p^{2} = p$. Therefore:
$|G/Z(G)|=p \implies G/Z(G)$ is cyclic $\implies G$ is abelian. But $G$ is non-abelian.
Now $|Z(G)|\neq 1$ also because $G$ is a $p-group$ and $p-groups$ have non-trivial center.
So, it must be that $|Z(G)|=p$.