I am doing some python practice and wrote a square code encoder/decoder. As described in drill #4 of Programming Practice Problems, the characters of a message are written left-to-right on a square grid, and each line of the encoded message is read by going top to bottom.

I'd like some feedback on what I wrote to see if there is anything I could have done better:

import math

def encode(message):
    message = message.replace(" ", "")
    square = math.ceil(math.sqrt(len(message)))

    # Create the "grid"
    rows = []
    for row in range(square):
        pos = row * square
        code = message[pos:pos + square]

    print("The grid:\n{}".format("\n".join(rows)))

    # Encode the message
    coded_msg = ""
    for col in range(square):
        for row in range(square):
            # The last column of the last row will throw an out of range error, so handle that
            if col < len(rows[row]):
                coded_msg += rows[row][col]
        coded_msg += " "

    return coded_msg.rstrip()

def decode(message):
    square = math.sqrt(len(message))
    col_count = math.ceil(square)
    row_count = math.ceil(len(message) / col_count)

    cols = message.split(" ")

    decoded_msg = ""
    for row in range(row_count):
        for col in cols:
            if row < len(col):
                decoded_msg += col[row]

    return decoded_msg

if __name__ == "__main__":

    print("Gimme a message:")
    in_message = input()

    coded_message = encode(in_message)
    print("--\nSquare message:\n{}".format(coded_message))
    decoded_message = decode(coded_message)
    #print("--\nDecoded message:\n{}".format(decoded_message))

    print("--\ninput: {}\noutput: {}".format(in_message, decoded_message))

2 Answers 2


In both the encoder and decoder, you append one character at a time to the resulting string using +=. That is a bad habit for performance: since Python strings are immutable, appending a character means allocating a new string, then copying the old contents of the string plus the suffix.

These solutions can be vastly simplified using Pythonic iteration techniques.

The first trick is to note that Python strings are sequences, and therefore support slicing using the s[i:j:k] operation to extract every kth character:

The slice of s from i to j with step k is defined as the sequence of items with index x = i + n*k such that 0 <= n < (j-i)/k. In other words, the indices are i, i+k, i+2*k, i+3*k and so on, stopping when j is reached (but never including j). When k is positive, i and j are reduced to len(s) if they are greater. … If i or j are omitted or None, they become “end” values (which end depends on the sign of k).

Secondly, whenever you see a pattern like:

array = []
for … in …:

… then that is an indication that you might want to use either a list comprehension or a generator expression, to define the whole array in just one statement.

In this case, you want to join a bunch of words using a space character as a delimiter. So, you need to write it as ' '.join(some generator expression producing the columns).

import math

def encode(plaintext):
    plaintext = plaintext.replace(" ", "")
    size = math.ceil(math.sqrt(len(plaintext)))
    return ' '.join(plaintext[i::size] for i in range(size))

To implement the decoder, you don't need to do any arithmetic. The trick here is to use itertools.zip_longest(*cols) to transpose the matrix and give you the rows.

Again, I recommend using ''.join(…) with some generator expressions.

from itertools import zip_longest

def decode(ciphertext):
    cols = ciphertext.split(" ")
    return ''.join(''.join(c for c in row if c) for row in zip_longest(*cols))
  • \$\begingroup\$ Wow, python is indeed beautiful! Fantastic answer, this teaches me a lot :) \$\endgroup\$
    – Spencer
    Apr 4, 2018 at 16:43
  • \$\begingroup\$ Is there a reason you used "len(plaintext) ** 0.5" to get the square instead of math.sqrt(plaintext)? And sorry for the noob question but what exactly does the "*" in "*cols" do? \$\endgroup\$
    – Spencer
    Apr 4, 2018 at 17:02
  • \$\begingroup\$ No reason for changing math.sqrt(); I've reverted it. The "splat" operator treats each element of the cols list as a separate parameter to zip_longest(). \$\endgroup\$ Apr 4, 2018 at 17:17

Couple of pointers.

In general it's good practice to avoid putting anything other than a call to another function (such as main()) in the if __name__ == "__main__" function, just to aid in ease of unit testing.

Document strings for functions are also nice to have:

def foo():
    '''This function does foo-all, always returns None'''

A nice (ish) one liner for the decode function might be something like this:

def encode(message):
    message = message.replace(" ", "")
    length = math.ceil(math.sqrt(len(message)))
    ''.join(message[i::length] i in range(length))

Although I admit it needs a bit of refactoring as its a tad long for one line (EDIT adapted a very nice improvement from 200_success 's answer, which they explain well). It does demonstrate that you can skip building the grid though.

  • \$\begingroup\$ Thanks Drgabble! Correct me if I'm wrong, but I thought it was "cleaner" to put any "testing" code in the main block so that the functions "encode" and "decode" can just be run from another module without worrying about extra junk at runtime? \$\endgroup\$
    – Spencer
    Apr 4, 2018 at 16:40

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.