I asked this question last night, but I have made some additional progress where I have another question.
Let $A$ be an $n \times n$ diagonalizable matrix with distinct eigenvalues $\lambda_1, \ldots, \lambda_k$ with corresponding multiplicities $d_1, \ldots, d_k$. Show that $$\dim(U) = d_1^2 + d_2^2 + \cdots + d_k^2,$$ where $U = \{B \in M_n \,|\, AB = BA \}$.
PROOF: Let $A$ be an $n \times n$ diagonalizable matrix with eigenvalues $\lambda_1, \ldots, \lambda_k$ and respective multiplicities $d_1, \ldots, d_k$. Since $A$ is diagonalizable, $D = P^{-1}AP$ is a diagonal matrix.
Observe that for any $B \in U$, $$P^{-1}AP = D \implies P^{-1}ABP = D(P^{-1}BP)$$ and likewise $$P^{-1}BAP = (P^{-1}BP)D.$$ This implies that we need only determine matrices $C$ such that $CD = DC$.
Let $$D = \begin{pmatrix}\lambda_1I_1 & & & \\ & \lambda_2I_2 & & \\ & & \ddots & \\ & & & \lambda_kI_k\end{pmatrix}.$$ Note that $I_i$ is the identity matrix of size $d_i \times d_i$. Partition $C$ as such: $$C = \begin{pmatrix}C_{11} & C_{12} & \cdots & C_{1k} \\ C_{21} & C_{22} & & C_{2k} \\ \vdots & \vdots & \ddots & \vdots \\ C_{k1} & C_{k2} & \cdots & C_{kk}\end{pmatrix}$$ Because $C$ and $D$ should commute, $CD = DC$ and so for any $i, j = 1, \ldots, k$, we have $$\lambda_iC_{ij} = C_{ij}\lambda_j \implies (\lambda_i - \lambda_j)C_{ij} = O.$$
Because each eigenvalue is distinct, it is not possible that $\lambda_i - \lambda_j = 0$ (for $i = j$), which leaves that $C_{ij} = O$ where $i \ne j$. So $C$ is block diagonal.
Here's where I'm stuck. I'm not sure what comes next to deduce $\dim(U)$. I think I need to find the dimension of each $C_{ii}$. How can I find the dimension of each $C_{ii}$? Surely it isn't as easy as saying that each $C_{ii}$ is $d_i \times d_i$, is it?