Highlight an equation within an align environment with color option

You can just add one more parameter to the macro. Below I have made the first parameter to specify the color -- it is optional and defaults to yellow if not provided as in the first example:

Notes:

• The main portion of this code is form Highlight an equation within an align environment, so please up vote that answer instead -- my addition here is just the addition of the first optional parameter.

Code:

\documentclass{article}
\usepackage{calc}
\usepackage{amsmath}
\usepackage{xcolor}

\newlength\dlf
\newcommand\alignedbox[3][yellow]{
  % #1 = color (optional, defaults to yellow)
  % #2 = before alignment
  % #3 = after alignment
  &
  \begingroup
  \settowidth\dlf{$\displaystyle #2$}
  \addtolength\dlf{\fboxsep+\fboxrule}
  \hspace{-\dlf}
  \fcolorbox{red}{#1}{$\displaystyle #2 #3$}
  \endgroup
}

\begin{document}
\begin{align*}
    \alignedbox{a}{=b} &                        c  &= d \\
                c &= d & \alignedbox[magenta!20]{i}{=k} \\
                e &= f &                        g  &= h  
\end{align*}
\end{document}

Another possibility is to exploit the hf-tikz package.

A couple of examples:

\documentclass[11pt,a4paper]{article}
\usepackage{amsmath}
\usepackage[customcolors,shade]{hf-tikz}

\begin{document}
\begin{align*}
     \tikzmarkin[fill=yellow]{first eq}a &= b\tikzmarkend{first eq} &          c  &= d \\
      c &= d &         i &= k \\
     \tikzmarkin{second eq}e &= f\tikzmarkend{second eq} &          \hfsetbordercolor{blue}\tikzmarkin[top color=white, bottom color=blue!20]{third eq}g  &= h\tikzmarkend{third eq}  
\end{align*}
\end{document}

which give the following result: