Highlighting pseudocode in algorithm

Question

Edit: I "solved" my main issues now. It's a bit ugly/hacky, but at least the result looks ok. With this edit I'm replacing the now obsolete old code with what I have now. But suggestions for improvement are always welcome.

Goal: Pseudocode with highlighted regions and corresponding annotations on the right hand side. How it looks now: Questions:

• The horizontal width of the colorboxes. I would want them to go all the way to the right and include the right hand side comments. (They should all end at the same spot.) "Solved" by duplicating my own boxit command. Ugly and not portable.
• The color is in front of the text, it would be more readable if it was in the background. Solved

The code for this is here:

\documentclass[a4paper,english]{article}
\usepackage[margin=3cm]{geometry}
% colors
\usepackage{xcolor, amsmath}
\definecolor{amaranth}{rgb}{0.9, 0.17, 0.31}
\colorlet{green}{green!20}
\colorlet{yellow}{yellow!60}
\colorlet{red}{red!30}
% pseudocode formatting
\usepackage{setspace}
\usepackage[linesnumbered,ruled,vlined]{algorithm2e}
\renewcommand{\KwSty}[1]{\textnormal{\textcolor{amaranth!90!black}{\bfseries #1}}\unskip}
\renewcommand{\ArgSty}[1]{\textnormal{\ttfamily #1}\unskip}
\SetKwComment{Comment}{\color{green!80!black!190}// }{}
\renewcommand{\CommentSty}[1]{\textnormal{\ttfamily\color{green!80!black!190}#1}\unskip}
\newcommand{\var}{\texttt}
\newcommand{\FuncCall}[2]{\texttt{\bfseries #1(#2)}}
\SetKwProg{Function}{function}{}{}
\SetKw{Continue}{continue}
\renewcommand{\ProgSty}[1]{\texttt{\bfseries #1}}
\DontPrintSemicolon
\SetAlFnt{\small}
\SetAlgorithmName{Pseudocode}{algorithmautorefname}
% pseudocode highlighting
\usepackage{tikz}
\usetikzlibrary{fit,calc}
% two slightly different boxit commands, to ensure the inner boxes end at the same spot
\newcommand{\boxit}[2]{
    \tikz[remember picture,overlay] \node (A) {};\ignorespaces
    \tikz[remember picture,overlay]{\node[yshift=3pt,fill=#1,opacity=.25,fit={($(A)+(0,0.15\baselineskip)$)($(A)+(.9\linewidth,-{#2}\baselineskip - 0.25\baselineskip)$)}] {};}\ignorespaces
}
\newcommand{\boxitt}[2]{
    \tikz[remember picture,overlay] \node (A) {};\ignorespaces
    \tikz[remember picture,overlay]{\node[yshift=3pt,fill=#1,opacity=.25,fit={($(A)+(0,0.15\baselineskip)$)($(A)+(.858\linewidth,-{#2}\baselineskip - 0.25\baselineskip)$)}] {};}\ignorespaces
}
\usepackage{algpseudocode}
\begin{document}
\begin{algorithm}
  \caption{find_best_split}
  \Function{find_best_split(X, gradients, curr_depth)}{
\tcp{determine node privacy budget}
\boxit{green}{6.5}
\eIf (\Comment{params.use\_decay}) {use\_decay} {  
    \boxitt{yellow}{3}
    \eIf (\Comment{curr\_depth == 0}) {curr\_depth == 0} {
        $\var{node\_budget} = \frac{\var{tree\_budget}}{2}$\;
    }{ 
        $\var{node\_budget} = \frac{\var{tree\_budget}}{4}$\;
    }
}{
    $\var{node\_budget} = \frac{\var{tree\_budget}}{6}$\;
}
\tcp{iterate over all possible splits} \boxit{green}{5}
\For(\Comment{number of cols in X}){feature\_index : features} {
    \boxitt{red}{3.5}
    \For(\Comment{number of rows in X}){feature\_value : X} {
        \If {use\_dp} {
            \Continue
        }
        $\var{gain} = \texttt{\emph{compute\_gain(X, gradients)}}$\;
    }
}
$\texttt{TreeNode *node = new TreeNode()}$\;
\Return{\var{node}}\;

}
\end{algorithm}
\end{document}