For a group $G$ such that $|G| = p^3$, for $p$ prime, either $|Z(G)| = p$ or $G$ is abelian. 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$.

The class equation you wrote is incorrect. You have to choose one element from each conjugacy class; if we denote by $g_1,\dots,g_k$ the ones chosen from the non singletons, we have
$$
|G|=|Z(G)|+\sum_{i=1}^k\frac{|G|}{\lvert\operatorname{Cent}(g_i)\rvert}
$$
This shows $|Z(G)|>1$, because $\operatorname{Cent}(g_i)\subsetneq G$, so the summation gives a multiple of $p$. Hence also $|Z(G)|$ must be a multiple of $p$. (Note that this just requires $|G|$ is a power of $p$.)
Thus only three cases are possible: $|Z(G)|=p$, $|Z(G)|=p^2$ or $|Z(G)|=p^3$. We need to exclude the second case.
Suppose $|Z(G)|=p^2$. Then $G/Z(G)$ has order $p$, so it's cyclic. If $gZ(G)$ is a generator, then for every $x\in G$ we have $xZ(G)=g^mZ(G)$, but this easily implies $G$ is abelian: contradiction.
I'm not sure your argument with the class equation is good.
