Enumerate numbers of the form a + b sqrt 2 - EPI

Enumerate numbers of the form a + b sqrt 2

GeneratingABSqrt2.java

  public static class ABSqrt2 implements Comparable<ABSqrt2> {
    public int a, b;
    public double val;

    public ABSqrt2(int a, int b) {
      this.a = a;
      this.b = b;
      val = a + b * Math.sqrt(2);
    }

    // clang-format off
    @Override
    public int compareTo(ABSqrt2 o) { return Double.compare(val, o.val); }
    // clang-format on
  }

  public static List<ABSqrt2> generateFirstKABSqrt2(int k) {
    SortedSet<ABSqrt2> candidates = new TreeSet<>();
    // Initial for 0 + 0 * sqrt(2).
    candidates.add(new ABSqrt2(0, 0));

    List<ABSqrt2> result = new ArrayList<>();
    while (result.size() < k) {
      ABSqrt2 nextSmallest = candidates.first();
      result.add(nextSmallest);

      // Add the next two numbers derived from nextSmallest.
      candidates.add(new ABSqrt2(nextSmallest.a + 1, nextSmallest.b));
      candidates.add(new ABSqrt2(nextSmallest.a, nextSmallest.b + 1));
      candidates.remove(nextSmallest);
    }
    return result;
  }

GeneratingABSqrt2Improved.java

  public static List<ABSqrt2> generateFirstKABSqrt2(int k) {
    // Will store the first k numbers of the form a + b sqrt(2).
    List<ABSqrt2> result = new ArrayList<>();
    result.add(new ABSqrt2(0, 0));
    int i = 0, j = 0;
    for (int n = 1; n < k; ++n) {
      ABSqrt2 resultIPlus1 = new ABSqrt2(result.get(i).a + 1, result.get(i).b);
      ABSqrt2 resultJPlusSqrt2
          = new ABSqrt2(result.get(j).a, result.get(j).b + 1);
      result.add(resultIPlus1.val < resultJPlusSqrt2.val ? resultIPlus1
                                                         : resultJPlusSqrt2);
      if (resultIPlus1.compareTo(result.get(result.size() - 1)) == 0) {
        ++i;
      }
      if (resultJPlusSqrt2.compareTo(result.get(result.size() - 1)) == 0) {
        ++j;
      }
    }
    return result;
  }