trjhtr

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]
 rows.append(code)

 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))

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 k^th 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 …:
 array.append(…)

… 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))

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'''
 pass

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.