How can I improve my spiral index function?

Question

Introduction

I am taking an online Introduction to Computer Science course for which I was given the task to create a function that takes in a coordinate and returns its corresponding "spiral index". I have a working function, but the online learning platform tells me that my code takes too long to execute (I am currently learning about code complexity and Big O notation).

The Problem Statement

All numbers on an infinite grid that extends in all four directions can be identified with a single number in the following manner:

Where 0 corresponds with the coordinate (0, 0), 5 corresponds with the coordinate (-1, 0), and 29 with (3, 2).

Create a function which returns the spiral index of any pair of coordinates that are input by the user.

Examples:

• spiral_index(10, 10) returns 380.
• spiral_index(10, -10) returns 440.
• spiral_index(3, 15) returns 882.
• spiral_index(1000, 1000) returns 3998000.

My Approach

def spiral_index(x, y):
    if x == 0 and y == 0:
        return 0
    pos = [1, 0]
    num = 1
    ring_up = 0
    ring_left = 0
    ring_down = 0
    ring_right = 0
    base = 3
    while pos != [x, y]:
        if ring_up < base - 2:
            pos[1] += 1
            ring_up += 1
            num += 1
        elif ring_left < base - 1:
            pos[0] -= 1
            ring_left += 1
            num += 1
        elif ring_down < base - 1:
            pos[1] -= 1
            ring_down += 1
            num += 1
        elif ring_right < base:
            pos[0] += 1
            ring_right += 1
            num += 1
        else:
            base = base + 2
            ring_up = 0
            ring_left = 0
            ring_down = 0
            ring_right = 0
    return num

The above code is able to find the correct index for every coordinate, and (on my laptop) computes spiral_index(1000, 1000) in just over 2 seconds (2.06).

Question

I have seen some solutions people have posted to a similar problem. However, I am wondering what is making my code so slow? To my knowledge, I believe that it is executing the function in linear time (is that right?). How can I improve the speed of my function? Can you post a faster function?

The course told me that a for loop is generally faster than a while loop, and I am guessing that the conditional statements are slowing the function down as well.

Any help on the matter is greatly appreciated!

Thanks in advance.

Answer 1

First of all, your solution takes two coordinates, A and B, and is O(AB), which can be considered quadratic. Second of all, the similar problem involves constructing the entire spiral, which means there is no better solution. You only have to index it, which means there's no reason to traverse the entire spiral. You can simply find which ring of the spiral you're on, then do some math to figure out the number. The center has 1 element, the next ring has 8, the next 16, the next 24, and it always increases by 8 from there. This solution is constant time and can almost instantly calculate spiral_index(1000, 1000).

def spiral_index(x, y):
    ax = abs(x)
    ay = abs(y)
    # find loop number in spiral
    loop = max(ax, ay)
    # one less than the edge length of the current loop
    edgelen = 2 * loop
    # the numbers in the inner loops
    prev = (edgelen - 1) ** 2
    if x == loop and y > -loop:
        # right edge
        return prev + y - (-loop + 1)
    if y == loop:
        # top edge
        return prev + loop - x + edgelen - 1
    if x == -loop:
        # left edge
        return prev + loop - y + 2 * edgelen - 1
    if y == -loop:
        # bottom edge
        return prev + x + loop + 3 * edgelen - 1
    raise Exception("this should never happen")