First step is to find a way to quickly calculate the largest square possible, this is not very difficult.
After this, I approached this problem in a similar way as I did for problem 199.
I start with the point (1, 0) and I generate it's maximum width using the above, which is ~0.61, then I generate the next 2 rectangle which have bottom left point (1.61, 0) and (1, 0.61).
The next problem is finding out when to stop generating new squares.
I'm sure there is a smarter way, but I just set a limit to how small the width can be and made it smaller until I found the correct answer.
To implement all of this while keeping track of the index I implemented a stack (I also tried recursion but it was more difficult and slower). I outline the first case below
Elements in the stack look like this [a, b, left, under, min_x]
(a, b) is the bottom left point of a square
left, under denotes the number of squares to the left and under
min_x is the limit for width as I explained above
Start with [1, 0, 0, 0, min_x]
a = 1, b = 0, left = 0, under = 0
find the width, ~0.61, using above formula and given point
if the width is greater than min_x
add [1 + 0.61, 0, 1, 0, min_x] notice we increased under by 1 because this square will be above the first square
add [1, 0.61, 0, 1, min_x] notice we increased left by 1 because this square will be to the left of the first square
After all this just loop through a sorted list and find the desired answer
Enter 2 numbers which denote the index (a, b) Note: if a, b ≥ 3, the code will give wrong answer
Code will output all n for which the index of Sn is (a, b)