# Create a snail matrix

Statement:

Create a script that, given an integer $n$, create a square matrix of dimensions $n \times n$ with the numbers from $1$ to $n^2$, arranged in a snail pattern.

Example:

1 2 3
8 9 4
7 6 5


I'm concerned by code quality, maintainability and general advices.

My implementation:

import collections.abc
import itertools
import math

type Matrix = list[list[int]]

def update_layer_edges(
c: collections.abc.Iterator, matrix: Matrix, layer_start: int, layer_end: int
) -> None:
"""Define the current layer tiles."""
for y, x in itertools.chain(
# top edge, where x increases and y is fixed to minimum
((layer_start, x) for x in range(layer_start, layer_end)),
# right edge where y increases and x is fixed to maximum
((y, layer_end - 1) for y in range(layer_start + 1, layer_end)),
# bot edge where x decreases and y is fixed to maximum
((layer_end - 1, x) for x in range(layer_end - 2, layer_start - 1, -1)),
# left edge where y decreases and x is fixed to minimum
((y, layer_start) for y in range(layer_end - 2, layer_start, -1)),
):
matrix[y][x] = next(c)

def square(size: int) -> list[list[int]]:
"""Create an size x size matrix arranged in a snail pattern."""
matrix = [[0] * size for _ in range(size)]
c = itertools.count(1)
layer_start = 0
layer_end = size
for _ in range(math.ceil(size / 2)):  # Repeat for each layer of the matrix
update_layer_edges(c, matrix, layer_start, layer_end)
layer_start, layer_end = layer_start + 1, layer_end - 1
return matrix



Test the implementation:

if __name__ == "__main__":
def print_matrix(matrix: Matrix) -> None:
for row in matrix:
print(" ".join(f"{num:>{padding}}" for num in row))

print_matrix(square(10))

assert square(1) == [[1]]
assert square(2) == [[1, 2], [4, 3]]
assert square(3) == [[1, 2, 3], [8, 9, 4], [7, 6, 5]]
assert square(4) == [[1, 2, 3, 4], [12, 13, 14, 5], [11, 16, 15, 6], [10, 9, 8, 7]]
assert square(5) == [
[1, 2, 3, 4, 5],
[16, 17, 18, 19, 6],
[15, 24, 25, 20, 7],
[14, 23, 22, 21, 8],
[13, 12, 11, 10, 9],
]


Visual representation of the matrix with code that generate the picture (I use matplotlib==3.8.4):

def show_snail(n) -> None:
import matplotlib.pyplot as plt

matrix = square(n)
# Using matplotlib to create a heatmap of the matrix
plt.figure(figsize=(8, 8))
plt.imshow(
matrix, cmap="viridis"
)  # You can choose any colormap like 'viridis', 'hot', 'cool', etc.
plt.gca().set_position(
[0, 0, 1, 1]
)  # Set the axes to fill the whole figure area
plt.gcf().set_size_inches(w=1, h=1)
plt.axis("off")  # Hide the axes
plt.show()

show_snail(5)


• You've got some good answers, but there's a fundamental issue with your program which might be necessary to address. The task said "PRINT the matrix", while your code first GENERATES the matrix and then prints it. Therefore, your algorithm has memory requirements of O(n^2). It should be possible to come up with an algorithm that will print the matrix without generating it and not require any extra memory.
– IMil
Apr 23 at 5:42
• @IMil Hum, I agree. The statement was written by me to summarize the task, but in fact, printing is not the goal. I'll edit the statement. I'll make a new version where the task is to print and work on an improvement. Thanks. Apr 23 at 7:42

It is great that you provided example code for us to run. The code layout is good and the implementation is compact.

Add a docstring to the top of the code to describe its purpose. For example, the "statement" from your question could be included in the docstring. Also, when I see the word "snail", I think "slow". You should mention that in this context, snail refers to a spiral pattern.

I'm not sure what you mean by a "layer", as in the update_layer_edges function name. I think you should add some details to the function docstring explaining what "layer" means. This will be helpful to the future person maintaining your code (even if it is you one year from now).

You should use the user-defined Matrix type in the type hint for the square function, as you have used it elsewhere. Change:

def square(size: int) -> list[list[int]]:


to:

def square(size: int) -> Matrix:


I think the following multiple assignment statement is easier to understand as two assignments on separate lines. Change:

    layer_start, layer_end = layer_start + 1, layer_end - 1


to:

    layer_start = layer_start + 1
layer_end   = layer_end   - 1


Consider changing the name of the square function because it is a bit generic. Perhaps change it to create_square to describe its action or what it returns.

• What do you think about adding this sentence in the renamed function update_layer ? "A layer is a set of tiles that are on the same distance from the center of the matrix." Apr 21 at 11:39
• I agree about the tuple increment, I considered it, and your opinion convinced me :) Apr 21 at 11:40
• @DorianTurba: That is a good description. Anything that helps someone (including the future you) understand the code better. Apr 21 at 11:40
• I renamed the function square to create_spiral_matrix. Apr 21 at 11:43

First of all, I wish there was more code in wild written like this. This implementation is self-contained and easy to read - I can quickly tell what it is doing just looking through the snippet.

Consider making your function more generic. You could parametrize Matrix by the contained type:

from collections.abc import Iterator

type Matrix[T] = list[list[T]]

def update_layer_edges[T](
c: Iterator[T], matrix: Matrix[T], layer_start: int, layer_end: int
) -> None:
...


This function does not really depend on the nature of elements, it would work exactly the same way if I wanted to pass, say, map(str, itertools.count(1)) as c.

I agree that both function and c argument should be renamed to something more understandable, but that's already covered in other answers.

And, as a minor stylistic nitpick, collections.abc and typing for hinting usually read better with import from. See e.g. Google style guide.

• I love your suggestion about T, it would make a great "follow-up" exercise asking spiral patterns of many things, strings, floats, pictures of snails ;) . About style, at first I was pro-google style, but since a few years, I tend to disagree on a few points, one of which is the import style. Since I do a lot of MR, review, teaching, mentoring and such, I share many snippets and top file imports may not be included. having the complete import namespace make reproduction and research faster and easier. It also make clear what's in the global namespace of the file, and what's not. Apr 22 at 11:30
• @DorianTurba yep, the styleguide is not perfect - but specifically this is perfectly fine IMO. We all know what Iterable or ParamSpec is, even without seeing the import - unless mentioned otherwise, such standard, often used standard library classes are recognizable. Reading a function signature with a couple of collections.abc.Iterable fully qualified mentions is much more annoying than possible ambiguity. And this lets pyupgrade (or ruff's UP codes) do their job better: it's a matter of 2 diff lines vs potentially hundreds of them. Apr 22 at 11:39
• And for a quick repro, from typing import * works just as well as import typing:) Apr 22 at 11:42
• I heard your point :) Apr 22 at 15:18

update_layer_edges isn't very accurate since a layer doesn't have "edges". I renamed it update_layer.

c could benefit from better typehinting and a better name:

- c: collections.abc.Iterator
+ n: collections.abc.Iterator[int]


A rework of update_layer_edges could reduce the number of parameters and the outside required value management. I also updated the docstring for the better. Variables layer_start and layer_end need to be renamed to lowest_idx and highest_idx is the entire function code.

- def update_layer_edges(
-     c: collections.abc.Iterator, matrix: Matrix, layer_start: int, layer_end: int
- ) -> None:
-     """Define the current layer tiles."""
-     for y, x in itertools.chain(
...
-     c = itertools.count(1)
-     layer_start = 0
-     layer_end = size
-     for _ in range(math.ceil(size / 2)):  # Repeat for each layer of the matrix
-         update_layer_edges(c, matrix, layer_start, layer_end)
-         layer_start, layer_end = layer_start + 1, layer_end - 1
-     return matrix
-------------------------------------------------------------------------------------
+ def update_layer(
+     n: collections.abc.Iterator[int], matrix: Matrix, layer: int
+ ) -> None:
+     """
+     Define the current layer tiles.
+
+     A layer is a set of tiles that are on the same distance from the border of the
+     matrix.
+
+     :param n: The generator of numbers that need to be placed in the matrix.
+     :param matrix: Target matrix to fill.
+     :param layer: Current layer to fill. 0 is the outermost layer.
+     """
+     lowest_idx = layer
+     highest_idx = len(matrix) - layer
...
+     n = itertools.count(1)
+     for layer in range(math.ceil(size / 2)):  # Repeat for each layer of the matrix
+         update_layer(n, matrix, layer)
+     return matrix


Regarding the plotting code:

1. Instead of commenting that cmap can be changed, parametrize it:

def show_snail(n: int, cmap: str = "viridis") -> None:
...
plt.imshow(matrix, cmap=cmap)
...

2. Instead of setting figsize((8, 8)) and then resizing with set_size_inches(1, 1), just set it once from the start:

plt.figure(figsize=(1, 1))

3. Instead of the pyplot interface which is convenient for exploratory code, use the native interface for production(ish) code:

from typing import Tuple
import matplotlib.pyplot as plt

def show_snail(n: int, cmap: str = "viridis") -> Tuple[plt.Figure, plt.Axes]:
matrix = square(n)

fig, ax = plt.subplots(figsize=(1, 1))
ax.imshow(matrix, cmap=cmap)
ax.set_position([0, 0, 1, 1])
ax.axis("off")
fig.show()

return fig, ax


Then the figure handles can be returned and reused:

snail_fig, snail_ax = show_snail(10, cmap="plasma")

# even if we create other intermediate plots
another_fig, another_ax = plt.subplots()
another_ax.plot(...)
...

# we easily still operate on the old snail figure