(a) Assume $A\in\mathbb{C}^{n\times n}$ has $n$ distinct eigenvalues. Prove that there are exactly $2^n$ distinct matrices $B$ such that $B^2 = A$ (i.e., in particular, there are no more than $2^n$ matrices with this property).
(b) How many such matrices $B\in\mathbb{C}^{3\times3}$ exist if $A=\begin{pmatrix}2&0&0\\0&2&0\\0&0&1\end{pmatrix}$? Why?
(a) It is clear that if $\lambda_i$ is an eigenvalue of $A$, then $\pm\sqrt\lambda_i$ is an eigenvalue of $B$, therefore $$B = \begin{pmatrix}\pm\sqrt\lambda_1&\cdots&0\\\vdots&\ddots&\vdots\\0&\cdots&\pm\sqrt\lambda_n\end{pmatrix}$$ satisfies $B^2=A$ and there are $2^n$ such matrices $B$ because of $2^n$ possibilities of rearranging different numbers of plus and minus signs on the diagonal.
(b) We can construct $2^3=8$ such matrices $B$ using approach of (a), but $B$ is not necessarily diagonal. Consider $$B=\begin{pmatrix}\sqrt2&-1&0\\0&-\sqrt2&0\\0&0&1\end{pmatrix}$$ which satisfies $B^2=A$ as well. So, there exist more than $2^3$ such matrices.
I can't give an explanation why there are no more than $2^n$ possibilities in the first case and more than $2^3$ in the second. It definitely follows from the fact that Jordan normal form is unique in (a) and isn't unique in (b) because of an eigenvalue with algebraic multiplicity $2$, however, I can't formulate it into a self-contained statement.