Conjugacy classes of a $p$-group

Let $K$ be a conjugacy class with more than one element. Since the order of $K$ divides $p^4$, it must be $p$, $p^2$ or $p^3$. Now $|K| = [G : C_G(g)]$, where $C_G(g)$ is the centralizer of some $g \in K$.

If $|K| = p^3$, then $|C_G(g)| = p$. This is not possible, since the center is always contained in the centralizer.

If $|K| = p^2$, then $|C_G(g)| = p^2$. Since the center is contained in $C_G(g)$, we get $C_G(g) = Z(G)$. Thus $g \in Z(G)$, implying $C_G(g) = G$ which is a contradiction.

Thus any conjugacy class with more than one element has exactly $p$ elements. Now use the class equation to find out the number of conjugacy classes in terms of $p$.