Here we expand a bit on Ilya Bogdanov's answer: $B = A + \Pi$, where $\Pi$ is the oblique projection matrix onto the null space of $A$ along the column space of $A$.
Oblique Projection Matrix
Given an $n \times n$ matrix $A$ with rank $r$. Compute its leading $r$ left singular vectors $\{ \vec{u}_i \in \mathbb{R}^n \mid 1 \le i \le r \}$. The oblique projection matrix onto the null space of $A$ or $\text{Null}(A)$ along the column space of $A$ or $\text{Col}(A)$ is:
$$
\Pi = I_n - \sum_{1 \le i \le r} \vec{u}_i \vec{u}_i^T \;. \tag{$\star$}
$$
Computational Cost of Constructing Oblique Projection Matrix
Note that computing this projection matrix only requires computing a compact SVD, i.e., finding the positive eigenvalues $\{ \lambda_i \}$ of $A^T A$ and their associated eigenvectors $\{ \vec{v}_i \}$. Then set
$$
\vec{u}_i = \frac{1}{\sqrt{\lambda_i}} A \vec{v}_i \tag{$\diamond$}
$$
for $1 \le i \le r$.
Why does ($\star$) work?
Recall that the left singular vectors are an orthonormal basis for $\text{Col}(A)$. Thus, one can always write the projection as:
$$
\Pi \vec{x} = \vec{x} -\sum_{1 \le i \le r} \alpha_i (\vec{u}_i \bullet \vec{x}) \vec{u}_i
$$
where the scalars $\{ \alpha_i \}$ are determined such that $\Pi \vec{x} \in \text{Null}(A) = \text{Null}(A^TA)$. In particular,
$$
(A A^T \vec{u}_j) \bullet ( \vec{x} - \sum_{1 \le i \le r} \alpha_i \vec{u}_i ) = 0 \implies \alpha_i = \vec{u}_i \bullet \vec{x} \quad \text{for $1 \le i \le r$}
$$
which gives the oblique projection map in ($\star$).
Transformed Matrix
Ilya proposed to the transform $B=A + \Pi$. This works because if $\vec{x}$ is an eigenvector of $A$ with nonzero eigenvalue then clearly $\vec{x} \in \text{Col}(A)$ and
$$
B \vec{x} = A \vec{x} = \lambda \vec{x}
$$
On the other hand, if $\vec{x}$ is an eigenvector of $A$ associated with a zero eigenvalue then $\vec{x} \in \text{Null}(A)$ or $\vec{x} \perp \text{Col}(A^T)$, and hence from ($\diamond$),
$$
B \vec{x} = \vec{x}
$$
So, $B$ fulfills the OP's desiderata.
MATLAB Implementation
One can replace $A$ below with any diagonalizable matrix.
% construction
A = [1 1 1; -2 -2 -1; 0 0 -1];
[n,n]=size(A);
r=rank(A);
[U,S,V]=svds(A,r);
B=A+(eye(n)-U*U');
% verification
[vectorsA,valuesA]=eig(A);
[vectorsB,valuesB]=eig(B);
valuesA=diag(valuesA);
ix=~(valuesA==0);
valuesA=valuesA(ix);
vectorsA=vectorsA(:,ix);
for i=1:length(valuesA)
B*vectorsA(:,i)-valuesA(i)*vectorsA(:,i)
end
