Encoding a monomial [lx,ly,lz] to contiguous index and decoding back

I want help or guidance to write one procedure called encode to compute index as in the following code snippet:

#include <iostream>

int encode(int lx, int ly, int lz) {
          return (ly+lz)*(ly+lz+3)/2 - ly;
        }


int main() {
  int L;
  // Replace with your desired value of L 
  for (L = 0; L <= 5; ++L) {
    int index = 0;
    for (int i = 0; i <= L; i++) {
      int lx = L - i;
      for (int j = 0; j <= i; j++, index++) {
        int ly = i - j;
        int lz = j;
        std::cout << "[" << lx << "," << ly << "," << lz << "] ->" << index <<" vs " << encode(lx, ly, lz)<< std::endl;
      }

    }
    std::cout << "-----------" << std::endl;
  }
  return 0;
}

'encode` above seems to work. I would like to implement a procedure named decode that calculates lx, ly, and lz based on the given index. Both encode and decode procedures are expected to be called frequently, making their performance crucial. Therefore, any suggestions to enhance their speed would be greatly appreciated.

As requested, I will describe how endode works.

The given procedure, encode(int lx, int ly, int lz), takes three non-negative integers lx, ly, and lz as input. These values represent a monomial (lx, ly, lz).

The procedure calculates an index value based on the monomial (lx, ly, lz). The index represents the index to store the information associated with the monomial in a contiguous array. The rules that follows the monomials are as follows:

1. The sum of the components lx, ly, and lz must be equal to L.
2. The value of L can take integer values from 0 to N, inclusive.
3. lx, ly, and lz must be non-negative integers.
4. L must also be a non-negative integer.

Initial Solution

As noted in the comments, the encode function can be replaced by return (ly+lz)*(ly+lz+3)/2 - ly;. This arises out of the fact that the encoding essentially follows triangular numbers: When we list the encodings in the desired index order, the list starts with one item with the maximum first lx value (L), and continues with two items with the next lower lx value (L−1), then three items with the next lower value and so on. So the number of items preceding the start of the current first index is the n^th triangular number, n(n+1)/2, where n is L−lx. That gives us the encoding for the first triple with the given lx value. From there, we want to account for the ly values. These are also listed in descending order, so the first one is L−lx. As they descend, we want to add more to the encoding, so we add L−lx−ly.

That gives us (L−lx)(L−lx+1)/2+L−lx−ly = (L−lx)(L−lx)/2-ly. L−lx is of course ly+lz, so this is (ly+lz)(ly+lz+3)−ly, and hence return (ly+lz)*(ly+lz+3)/2 - ly;.

To reverse this, we recognize that, at the first list item with a given value for lx, the L−lx−ly contribution is zero, and we have the triangular number equaling the encoding e, so (L−lx)(L−lx+1)/2 = e. Solving this quadratic equation for lx gives lx = L + 1½ - sqrt(2¼ + 2e). (The quadratic has two solutions, but this is the one of interest to us, as the other would make lx greater than L.)

Thus, given an encoding e, we can set lx = L + 1.5 - std::sqrt(2.25 + 2*e);. This is exact when e is the encoding of the first list item for the given lx. When e is a later encoding for the same lx with different ly and lz, the formula gives a higher value, but it does not reach the next integer until lx changes. So the truncation to integer in the assignment yields the correct lx for all encodings with a given lx value.

I will leave the algebra for finding ly to the reader and give it as ly = (L-lx)*(L-lx+3)/2 - e;. And of course lz is merely lz = L-lx-ly;.

Thus a decode routine is:

#include <cmath>
void decode(int L, int e, int &lx, int &ly, int &lz)
{
	lx = L + 1.5 - std::sqrt(2.25 + 2*e);
	ly = (L-lx)*(L-lx+3)/2 - e;
	lz = L-lx-ly;
}

Caution

The above requires that sqrt give a correct result when the result is exactly an integer. In C++ implementations where sqrt may deliver a poorly computed result, it may be necessary to make an adjustment, either by testing lx and incrementing if necessary or adding an adjustment, as with L + 1.5 - std::sqrt(2.25 + 2*e) + 1e-10;. A reasonable value for such an adjustment could be computed based on a bound on L.

Update: A nice way to do this is to subtract ½ from e in the calculation of lx. When 2¼+2e would be a perfect square, this moves it away from that (in the desired direction, making it smaller so subtracting it from L+1½ produces a large value, which will be truncated to the proper integer), which avoids bad consequences of small errors in sqrt, but it does not push other values of e across the next boundary where the calculated value of lx would change. And, since e appears as 2e, we can do this by subtracting 1 from 2¼+2e, thus using 1¼+2e: lx = L + 1.5 - std::sqrt(1.25 + 2*e);.