2 files changed, 143 insertions, 0 deletions
diff --git a/robots/algorithms.pyx b/robots/algorithms.pyx
new file mode 100644
index 0000000..620cfce
--- /dev/null
+++ b/robots/algorithms.pyx
@@ -0,0 +1,135 @@
+# cython: language_level=3
+
+cimport cython
+
+import numpy as np
+cimport numpy as np
+
+ITYPE = np.int32
+ctypedef np.int32_t ITYPE_t
+
+ctypedef cython.bint bool
+
+from robots.constants import DIRECTIONS
+
+INFINITY = float('inf')
+FAKE_INFINITY = 1<<31 - 1
+
+def distance_relaxer(distances, iterations=None):
+    width = len(distances[0])
+    height = len(distances)
+
+    predecessors = [
+        [None] * width
+        for _ in range(height)
+    ]
+
+    convert = {
+        INFINITY: FAKE_INFINITY,
+        None: -1,
+    }
+    unconvert = {
+        FAKE_INFINITY: INFINITY,
+        -1: None,
+    }
+    _distances = np.array(
+        [
+            [
+                convert.get(v, v)
+                for v in row
+            ]
+            for row in distances
+        ],
+        dtype=ITYPE,
+        order='c',
+    )
+
+    if iterations is None:
+        iterations = -1
+
+    _distance_relaxer(
+        _distances,
+        predecessors,
+        iterations,
+        width, height,
+    )
+
+    for y in range(height):
+        for x in range(width):
+            v = _distances[y, x]
+            distances[y][x] = unconvert.get(v, v)
+
+    return predecessors
+
+@cython.boundscheck(False)
+cdef void _distance_relaxer(
+    np.ndarray[ITYPE_t, ndim=2, mode='c'] distances,
+    predecessors,
+    int iterations,
+    Py_ssize_t width,
+    Py_ssize_t height,
+    DIRECTIONS=DIRECTIONS,
+):
+    cdef Py_ssize_t x, y
+    cdef Py_ssize_t nx, ny
+    cdef ITYPE_t distance, this_distance
+
+    cdef int i = 0
+    cdef bool updated
+
+    while i != iterations:
+        updated = False
+
+        for y from 0 <= y < height:
+            for x from 0 <= x < width:
+                distance = distances[y, x]
+                if distance == -1:
+                    continue
+
+                nx = x
+                ny = (y - 1) % height
+                this_distance = distances[ny, nx]
+                if this_distance != -1:
+                    if this_distance != FAKE_INFINITY:
+                        this_distance += 1
+                    if this_distance < distances[y, x]:
+                        distances[y, x] = this_distance
+                        predecessors[y][x] = 'U'
+                        updated = True
+
+                nx = x
+                ny = (y + 1) % height
+                this_distance = distances[ny, nx]
+                if this_distance != -1:
+                    if this_distance != FAKE_INFINITY:
+                        this_distance += 1
+                    if this_distance < distances[y, x]:
+                        distances[y, x] = this_distance
+                        predecessors[y][x] = 'D'
+                        updated = True
+
+                nx = (x - 1) % width
+                ny = y
+                this_distance = distances[ny, nx]
+                if this_distance != -1:
+                    if this_distance != FAKE_INFINITY:
+                        this_distance += 1
+                    if this_distance < distances[y, x]:
+                        distances[y, x] = this_distance
+                        predecessors[y][x] = 'L'
+                        updated = True
+
+                nx = (x + 1) % width
+                ny = y
+                this_distance = distances[ny, nx]
+                if this_distance != -1:
+                    if this_distance != FAKE_INFINITY:
+                        this_distance += 1
+                    if this_distance < distances[y, x]:
+                        distances[y, x] = this_distance
+                        predecessors[y][x] = 'R'
+                        updated = True
+
+        if not updated:
+            break
+        i += 1
diff --git a/setup.py b/setup.py
new file mode 100644
index 0000000..77263a7
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,8 @@
+from setuptools import find_packages, setup
+from Cython.Build import cythonize
+
+setup(
+    name='robots',
+    packages=find_packages(),
+    ext_modules = cythonize('robots/*.pyx')
+)