Relation between matrix power and Jordan normal form 
(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.
 A: You showed that $A$ has at least $2^n$ square roots. Assume $B^2 = A$. Since $\sigma(B)^2 =\sigma(B^2) = \sigma(A)$, it follows that
$$\sigma(B) = \{\varepsilon_1\sqrt{\lambda_1}, \ldots, \varepsilon_n\sqrt{\lambda_n}\}$$
for some $\varepsilon_1, \ldots, \varepsilon_n \in \{-1,1\}$ where $\sigma(A) = \{\lambda_1, \ldots \lambda_n\}$.
Set $D = \operatorname{diag}(\varepsilon_1\sqrt{\lambda_1}, \ldots, \varepsilon_n\sqrt{\lambda_n})$. $A$ has $n$ distinct eigenvalues so it is diagonalizable, i.e. there exists an invertible matrix $P$ such that $P^{-1}AP = \operatorname{diag}(\lambda_1, \ldots, \lambda_n) = D^2$.
Assume that $C^2 = A$ and that $\sigma(C) = \sigma(B)$. $B$ and $C$ both have $n$ distinct eingenvalues so there exist invertible matrices $Q,R$ such that 
$$Q^{-1}BQ = D = R^{-1}CR$$
It follows 
$$Q^{-1}AQ = Q^{-1}B^2Q = (Q^{-1}BQ)^2 = D^2 = (R^{-1}CR)^2 = R^{-1}C^2R = R^{-1}AR$$
If $e_1, \ldots, e_n$ is the standard basis, it follows that $A(Pe_i) = \lambda_iPe_i, A(Qe_i) = \lambda_iQe_i, A(Re_i) = \lambda_iRe_i$.
Since $\lambda_1, \ldots, \lambda_n$ are all distinct, the eigenvectors are unique up to scalar multiplication so it follows that $Qe_i = \sigma_i Pe_i, Re_i = \pi_i Pe_i$ for some nonzero scalars $\sigma_i, \pi_i$, or $Q = P\Sigma, R = P\Pi$ for some invertible diagonal matrices $\Sigma, \Pi$.
Finally,
$$B = QDQ^{-1} = (P\Sigma)D(P\Sigma)^{-1} = P(\Sigma D\Sigma^{-1})P^{-1} =\\ PDP^{-1} = P(\Pi D\Pi^{-1})P^{-1} = (P\Pi)D(P\Pi)^{-1} = R^{-1}DR = C$$
Hence $B$ is uniquely determined by its spectrum $\sigma(B)$, which is uniquely defined by choosing $\varepsilon_i \in \{-1,1\}$ which can be done it $2^n$ ways. We conclude that there are $2^n$ such matrices $B$.
A: See my answer in 
Can we reduce finding matrix roots to finding roots of Jordan blocks?
In (b), your matrix $A$ has $2$ Jordan blocks associated to the same eigenvalue $2$. Then $A$ admits an infinity of square roots.
Here the infinity of matrices $B\in M_2(\mathbb{C})$ s.t. $tr(B)=0,\det(B)=-2$ are among the square roots of $2I_2$.
EDIT. I forgot to write that, in case (a), there are no more than $2^n$ square roots because, necessarily $AB=BA$.
