# Finding the dimension of $S = \{B \in M_n \,|\, AB = BA\}$, where $A$ is a diagonalizable matrix

I need some help proving the following:

Let $A \in M_{n \times n}$ be a diagonalizable matrix with distinct eigenvalues $\lambda_1, \ldots , \lambda_k$ and corresponding multiplicities $d_1, \ldots, d_k$. Given that $S = \{B \in M_n \,|\, AB = BA\}$, prove that $\dim (S) = d_1^2 + \cdots + d_k^2$.

Here's what I've done:

Since $A$ is diagonalizable, we have that $D = P^{-1}AP$, where $D$ is a diagonal matrix. We may then write, $$P^{-1}AP = \begin{bmatrix}\Lambda_1 & & & \\ & \Lambda_2 & & \\ & & \ddots & \\ & & & \Lambda_k\end{bmatrix} = D,$$ where $$\Lambda_i = \underbrace{\begin{bmatrix}\lambda_i & & & \\ & \lambda_i & & \\ & & \ddots & \\ & & & \lambda_i\end{bmatrix}}_{d_i \, \text {columns}}.$$

Obviously the size of each $\Lambda_i$ is $d_i \times d_i = d_i^2$, which I'm thinking we need to show that each block contributes this to the total dimension as we work along the diagonal, but I'm having a little trouble coming to this conclusion via the why and how.

Clearly each column vector of $D$ and consequently each column vector of each $\Lambda_i$ is linearly independent, but I am unable to see what comes next to receive the squared part. Can someone provide a hint as to how I can make this conclusion?

Hint: Write $D$ as a block matrix $$D= \pmatrix{ \lambda_1 I \\ & \lambda_2 I \\ && \ddots \\ &&& \lambda_k I }$$ Noting that $$P^{-1}[AB]P= D(P^{-1} BP), \qquad P^{-1}[BA]P= (P^{-1} BP)D$$ It suffices to determine which matrices $M$ satisfy $MD=DM$.
To that end, write $M$ as a block matrix (partitioned in the same way as $D$), and compute the products $DM$ and $MD$. Determine that all blocks of $M$ must be zero except those on the diagonal.
Hint. Let $B = [b_{ij}]$ be a $n\times n$ matrix which commutes with a diagonal matrix $D=\mbox{diag}(\lambda_1,\lambda_2,\dots,\lambda_n)$, that is $DB=BD$.
Then by computing $DB$ and $BD$, we get that for all $1\leq i,j\leq n$, $$\lambda_i b_{ij} = b_{ij}\lambda_j$$ which implies that $(\lambda_i - \lambda_j) b_{ij} = 0$. Hence either $\lambda_i = \lambda_j$ or $b_{ij} = 0$.