Compute the kth element in spiral order for an m x n 2D array in O(1) time (Python)

kth_element_spiral.py

# (Variant #6 for exercise 5.18 on EPI (Elements of Programming Interviews)) (September 2018 edition)
# Consider a 10x7 (mxn) matrix, in which we get 30 elements for the the outer ring, the next outer ring would have 22 elements, then 14 elements, and the most inner ring has the remaining elements. The number of elements per ring is given by 2 x (m - (1 + (2 x (r-1) ))) + 2 x (n - (1 + (2 x (r-1) ))), for the rth ring.
# Save from the most inner ring, the difference between the number of elements of each adjacent ring is 8 elements. If we want to know the number of elements of the current ring plus all the other previous ones, we get an arithmetic series (https://en.wikipedia.org/wiki/Arithmetic_progression#Sum).
# The sum of all the elements until a given r is given by sum = (r(a1 + ar)) / 2, being a1 = 2(m-1) + 2(n-1), ar = 2 x (m - (1 + (2 x (r-1) ))) + 2 x (n - (1 + (2 x (r-1) ))). If we solve the equation in relation to r, using the quadratic formula to disentangle the final polynomial, we reach r = math.floor((-(m + n) + math.sqrt((m+n)**2 - 4*k)) / (-4))
# Now we can know in which ring the kth element is, and also know the number of elements leading up to this ring. That is the base of the below algorithm.
# The rest is working with these two pieces of information to work out an offset and figure out where that offset falls in the ring (top, right, bottom or bottom portion)
# This algorithm has O(1) time complexity.
# *Note that the algorithm below computes the x and y on the array for a given k, m and n, which is the core part of this problem. Getting the array item for these coordinates is trivial from here.*

import math

def kth_element_spiral(k: int, m: int, n: int):
    first_ring_max = 2 * (m-1) + 2 * (n - 1)

    ring, ring_start_elements = None, None

    if (k < first_ring_max):
        ring = 0
        ring_start_elements = 0
    else:
        ring = int(math.floor((-(m + n) + math.sqrt((m+n)**2 - 4*k)) /(-4)))
        ring_start_elements = int((ring * (first_ring_max + 2 * (m-(1 + 2*(ring - 1))) + 2 * (n-(1 + 2*(ring - 1))))) / 2)

    offset = k - ring_start_elements

    width = (m - ring*2) - 1
    height = (n - ring*2) - 1

    if (offset <= width): # top
        x = ring + offset 
        y = ring
        return (x, y)
    elif (offset <= width + height): # right
        x = m - ring - 1
        y = ring + (offset - width)
        return (x, y)
    elif (offset <= width + height + width): # bottom
        x = width - (offset - width - height)
        y = n - ring - 1
        return (x, y)
    else: # left
        x = ring
        y = height - (offset - width - height - width)
        return (x, y)


# Some test cases    
for i in range(10*6):
    print(i, kth_element_spiral(i, 10, 6))

for i in range(7*5):
    print(i, kth_element_spiral(i, 7, 5))

for i in range(9*6):
    print(i, kth_element_spiral(i, 9, 6))