I got the above problem in an exam, the problem stated to show there is at most $2^n$ such matrices.
What I did was to show the matrix is similar to a matrix with a diagonal matrix with distinct complex numbers, and now if $A$ is such a matrix and $B$ is a square root. Then $B$ is diagonalizable so since $A$ and $B$ commute there is a matrix $P$ such that that $P^{-1}BP$ and $P^{-1}AP$ are both are diagonal.
So my argument was that to consider the diagonal case (w.l.g) and show that each diagonal entry of $P^{-1}BP$ should be root of the corresponding diganoal entry of $P^{-1}AP$. So we can have atmost 2 options for each and overall we get $2^n$.
Now the problem I just realized is that I just saw that for any invertible diagonalizable matrix $A^2$, $A$ is also diagonalizable. So then we can apply the same logic as what I just did to any such invertible matrix. But I also know the identity has infinitely many roots. So where did I go wrong?
Thank You
edit: added non zero