Determine the Last Digits of N^2 == N in Java


The challenge

Let’s call an integer number N ‘green’ if  ends with all of the digits of N. Some examples:

5 is green, because 5² = 25 and 25 ends with 5.

11 is not green, because 11² = 121 and 121 does not end with 11.

376 is green, because 376² = 141376 and 141376 ends with 376.

Your task is to write a function green that returns nth green number, starting with 1 – green (1) == 1

Data range

n <= 500

The solution in Java code

import java.math.BigInteger;
import java.util.TreeSet;

public class GreenNumbers {
    public static BigInteger get(int n) {
        if (n == 1) {
            return BigInteger.ONE;
        }

        TreeSet<BigInteger> treeSet = new TreeSet<>();
        addNumber(treeSet, n + n / 5);
        BigInteger[] bigIntegers = treeSet.toArray(new BigInteger[]{});
        return bigIntegers[n - 2];
    }

    public static void addNumber(TreeSet<BigInteger> treeSet, int number) {
        int digitNo = 2;
        int counter = 2;
        BigInteger num = new BigInteger("5");
        treeSet.add(num);
        BigInteger num2 = new BigInteger("6");
        treeSet.add(num2);
        while (counter < number) {
            BigInteger numSquared = num.multiply(num);
            BigInteger powOfTen = new BigInteger("10").pow(digitNo);
            num = numSquared.mod(powOfTen);
            num2 = powOfTen.add(new BigInteger("1")).subtract(num);
            treeSet.add(num);
            treeSet.add(num2);
            counter += 2;
            digitNo++;
        }
    }
}

Test cases to validate our solution

import org.junit.Test;
import static org.junit.Assert.assertEquals;
import org.junit.runners.JUnit4;
import java.math.BigInteger;


public class SolutionTest {
    @Test
    public void basicTests() {
        assertEquals(new BigInteger("1"),  GreenNumbers.get(1));
        assertEquals(new BigInteger("5"),  GreenNumbers.get(2));
        assertEquals(new BigInteger("6"),  GreenNumbers.get(3));
        assertEquals(new BigInteger("25"), GreenNumbers.get(4));
    }
    
    @Test
    public void biggerTests() {
        assertEquals(new BigInteger("2890625"), GreenNumbers.get(12));
        assertEquals(new BigInteger("7109376"), GreenNumbers.get(13));
    }
    
    @Test
    public void advancedTests() {
        assertEquals(new BigInteger("6188999442576576769103890995893380022607743740081787109376"),       GreenNumbers.get(100));
        assertEquals(new BigInteger("9580863811000557423423230896109004106619977392256259918212890625"), GreenNumbers.get(110));
    }
}