I've taken an approach and it really feels like it should be getting me there, but I can't seem to close it off. Any hints?
Attempt: If $G$ is a group of order $p^2$ then as a $p$-group $Z(G)$ is non-trivial. Because $Z(G)$ is a subgroup of $G$ Lagrange's theorem tells us that the order of $Z(G)$ must divide the order of $p$, and since $Z(G)$ is non-trivial we must have $|Z(G)| = p$ or $p^2$. If $|Z(G)|$ has order $p^2$ we are finished since then $Z(G) = G$ which implies $G$ is abelian. So assume that $|Z(G)| = p$, i.e $Z(G) = \langle g \rangle$ for some $g \in Z(G)$. Because $p$ is prime it follows that $Z(G)$ is a cyclic group. Moreover $Z(G) \unlhd G$ and so we can form the quotient group and by Lagrange's $$ |G/Z(G)| =\frac{|G|}{|Z(G)|} = p, $$ and so the quotient group is also cyclic and so $G/Z(G) = \langle hZ(G) \rangle$ for some $h \in G$.
I feel like there must be some way to finish this off right? I have yet to learn about products of groups but maybe that's how you're supposed to proceed?
Edit: It occurred to me that $G$ is partitioned into the cosets of $Z(G)$ in $G$ and so the proof can be finished.
Let $G$ be a group of order $p^2$. Then $Z(G)$ is a non-trivial normal subgroup of $Z(G)$. The case where $|Z(G)| = p^2$ is trivial, so assume that $|Z(G)| = p$. Then by Lagrange's theorem $$ |G/Z(G)| = \frac{|G|}{|Z(G)|} = p $$ and so as a group of prime order $G/Z(G)$ is cyclic. Let $g \in G$ be a representative such that $G/Z(G) = \langle gZ(G) \rangle$. Because $G$ is partitioned into it's cosets, any element of $G$ takes the form $g^\ell h$ since each coset is generated by $gZ(G)$! This means that for any $a,b \in G$ we have that $a = g^kh_1$ and $b = g^mh_2$ where $h_1,h_2 \in Z(G)$. Then \begin{align*} ab &= (g^kh_1)(g^mh_2) \\ &=h_1g^kg^mh_2 \\ & = h_1g^mg^kh_2 \\ & = (g^m h_2)(g^kh_1)\\ & = ba. \end{align*} These lines follow by the fact that $h_1,h_2 \in Z(G)$ and so commute with every element of $G$, and powers of $g \in G$ commute in the same way elements of a cyclic group do.
