# Vigenère Cipher problem in competitive programming

I've just started doing programming problems.

https://icpcarchive.ecs.baylor.edu/external/58/p5880.pdf

The code below works fine, but my run time is 2.715 seconds and the time limit is 3 seconds.

How can I improve speed of this code?

I've already started to use sys.stdin instead of input().

import string
import sys

while True:
key = int(key) if key.isdigit() else key
if key == 0: break

decrypted_msg = int(decrypted_msg) if decrypted_msg.isdigit() else decrypted_msg
if decrypted_msg == 0: break

key = (len(decrypted_msg)//len(key) + 1) * key
key = key[0:len(decrypted_msg)]

encrypted_msg = ""

for i in range(len(decrypted_msg)):
key_idx = string.ascii_uppercase.index(key[i])
deciphered_idx = string.ascii_uppercase.index(decrypted_msg[i])

enciphered_idx = (deciphered_idx + key_idx + 1) % len(string.ascii_uppercase)
encrypted_msg += string.ascii_uppercase[enciphered_idx]

print(encrypted_msg)

• Where can solutions be submitted? Commented Feb 25, 2021 at 23:08
• Welcome to Code Review! Thanks for providing a link to the specification, but do note that a question should be complete in itself. Please edit to include a short description of the requirements in the body of the question. Obviously, you should take care not to infringe copyrights in the process of fixing this. Commented Feb 26, 2021 at 16:37

• decrypted_msg is misleading. It's not decrypted but rather "not encrypted". But really just call it "plaintext", as the problem statement does.
• Similarly, "enciphered" and "encrypted" are not just inconsistent (make up your mind), but go against the problem statement's term "ciphertext".
• The end-signal string '0' can simply be checked as a string.
• No need to check the plaintext against '0'.
• Mapping str.strip over sys.stdin makes the code a bit nicer. We can even use a nice for statement then.
• Better just zip the plaintext and key instead of iterating over indices.
• Searching the letters in the alphabet string isn't super terrible, but it's apparently faster to look up letter pairs in a dictionary.
• itertools.cycle is a nicer (and I think faster) way to cycle the key.
• Repeatedly extending a string with += might degrade to quadratic time overall. Better use ''.join on an iterable.
• A little time can be saved by storing string.ascii_uppercase.index in a variable instead of looking up the attributes ascii_uppercase and index over and over again.

## Solution 1: My fastest[*], using a lookup-table mapping letter pairs to single letters

This gets accepted in 0.382 seconds here, where I believe you submitted yours (tried four times, impressively always got the exact same time).

[*] In my other answer benchmarking a worst case, actually my solution 3 is fastest.

import sys
from itertools import cycle
from string import ascii_uppercase as abc

table = {(abc[i], abc[j]): abc[(i + j + 1) % 26]
for i in range(26)
for j in range(26)}.get

lines = map(str.strip, sys.stdin)
for key in lines:
if key == '0':
break
plaintext = next(lines)

print(''.join(map(table, zip(plaintext, cycle(key)))))


I did get it down to 0.372 seconds by using sys.stdout.write, but meh, not worth it.

## Solution 2: My fastest with str.index:

Here's the fastest I got with str.index, got accepted in 0.699 seconds. Note that my bcd is twice the alphabet and starts with "B", so I can use bcd[p + k] instead of the slower abc[(p + k + 1) % 26].

from string import ascii_uppercase as abc
import sys
from itertools import cycle

index = abc.index
bcd = abc[1:] + abc

lines = map(str.strip, sys.stdin)
for key in lines:
if key == '0':
break
plaintext = next(lines)

key = cycle(map(index, key))
plaintext = map(index, plaintext)

print(''.join([bcd[p + k] for p, k in zip(plaintext, key)]))


## Solution 3: O(|key|) Python operations

Since the plaintext has up to 100,000 letters but the key has only up to 1,000 letters, and since Python operations are slow, I also tried doing just O(|key|) Python operations instead of O(|plaintext|) Python operations. By using str.translate. For example if the key is 'KEY', then the entire string slice plaintext[0::3] is translated with the key letter 'K'. And [1::3] and [2::3] get translated with 'E' and 'Y', respectively. Sadly, it was slower, got accepted in 0.985 seconds:

import sys
from string import ascii_uppercase as abc

tables = {c: str.maketrans(abc, abc[i:] + abc[:i])
for i, c in enumerate(abc, 1)}

lines = map(str.strip, sys.stdin)
for key in lines:
if key == '0':
break
plaintext = next(lines)

ciphertext = [None] * len(plaintext)
n = len(key)
for i, k in enumerate(key):
ciphertext[i::n] = plaintext[i::n].translate(tables[k])

print(''.join(ciphertext))


## Solution 3b: Using bytearray

Solution 3 but with bytes/bytearray instead of str. Accepted in 0.619 seconds. Not sure whether it's the input/output or the bytearray.translate that's responsible for the speed-up:

import sys
from string import ascii_uppercase as abc

abc = abc.encode()
tables = {c: bytearray.maketrans(abc, abc[i:] + abc[:i])
for i, c in enumerate(abc, 1)}

lines = map(bytes.strip, sys.stdin.buffer)
write = sys.stdout.buffer.write
for key in lines:
if key == b'0':
break
plaintext = next(lines)

ciphertext = bytearray(plaintext)
n = len(key)
for i, k in enumerate(key):
ciphertext[i::n] = ciphertext[i::n].translate(tables[k])

write(ciphertext)
write(b'\n')

• This is a good answer -- well done. I had a couple beers and typed my response without thinking too hard about the algorithm. Anyway, I encourage the OP to focus on this reply, but I'll leave my answer up, because it does have some useful points about how to set up an overall program in a useful way for experimentation.
– FMc
Commented Feb 26, 2021 at 2:37

Start with some of the low-hanging fruit. (1) Inside your primary for-loop, you repeatedly used index() to hunt for a value that can be easily computed in advance. (2) You are assembling the output data via string concatenation, but Python strings are immutable, which implies the creation of an entirely new string on each iteration of the loop. That can become problematic when dealing with larger inputs. Here's an edited version with those issues in mind. It sticks fairly closely to your original code.

import string
import sys

# Create an dict mapping each uppercase letter to its index in
# string.ascii_uppercase. In addition, (and these changse are much less
# important) make letters a local variable so we don't have to perform an
# attribute lookup on string.ascii_uppercase inside the loop. And define
# a constant to eliminate the need to call len() repeatedly.
letters = string.ascii_uppercase
letter_indexes = {c : i for i, c in enumerate(letters)}
n_letters = len(letters)

while True:

# Simplifying the loop-break logic.
if key == '0':
break

# Simplifying the key-creation logic.
# We don't care whether key is too long.
key = (len(decrypted_msg)//len(key) + 1) * key

# We will append characters to a list rather
# than creating a fresh string each time.
encrypted_chars = []

for i in range(len(decrypted_msg)):
key_idx = letter_indexes[key[i]]
deciphered_idx = letter_indexes[decrypted_msg[i]]

enciphered_idx = (deciphered_idx + key_idx + 1) % n_letters
encrypted_chars.append(letters[enciphered_idx])

# Assemble the final string.
enciphered_msg = ''.join(encrypted_chars)
print(enciphered_msg)


Your program is difficult to test and debug, due to its larger design. For example, your linked pdf provided some sample inputs and expected outputs, but the program itself did not provide easy ways to toggle between its "real life" usage (reading from stdin, in this case) and various testing/debugging usages. This adds friction to attempts to experiment with program, hampering the ability to make refinements, simplifications, and speed enhancements. A more flexible approach is to keep the Vigenère algorithm completely isolated from the orchestration tasks of handling inputs and different usage modes.

The code below provides a sketch of a more flexible structure. It also illustrates how to take advantage of zip() and itertools.cycle() to simplify the code logic -- and, in particular, to rely less on manual indexing into strings. Although this part of the answer was not intended to provide additional speed-ups, at least in my not-very-sophisticated experiments, these changes also made the algorithm notably faster.

Finally, the edited vigenere() function provides a nice illustration of the way that carefully chosen, shorter variable names can increase code readability and clarity -- provided that those shorter names are within a narrow context (i.e., a short function) and have sufficient contextual information (e.g., code comments and nearby function/variable names that are more explicit).

import string
import sys
from itertools import cycle

# Test data from the pdf file.
TESTS = [
['ICPC', 'THISISSECRETMESSAGE', 'CKYVRVIHLUUWVHIVJJU'],
['ACM', 'CENTRALEUROPEPROGRAMMINGCONTEST', 'DHAUUNMHHSRCFSEPJEBPZJQTDRAUHFU'],
['LONGKEY', 'CERC', 'OTFJ'],
]

def main(args):
if args and args[0] == 'test':
# Make sure our vigenere() is still OK.
for key, text, expected in TESTS:
encrypted = vigenere(text, key)
print(encrypted == expected, encrypted)
else:
# The coding challenge use case.
while True:
if key == '0':
break
encrypted = vigenere(text, key)
print(encrypted)

def vigenere(text, key):
# Takes text and a key.
# Returns encrypted text.

# Prepare letter index lookup.
letters = string.ascii_uppercase
indexes = {c : i for i, c in enumerate(letters)}
n_letters = len(letters)

# Encrypt the text.
encrypted_chars = []
for t, k in zip(text, cycle(key)):
ti = indexes[t]
ki = indexes[k]
ei = (ti + ki + 1) % n_letters
encrypted_chars.append(letters[ei])

# Return as string.
return ''.join(encrypted_chars)

if __name__ == '__main__':
main(sys.argv[1:])

• Good point with the separate vigenere function and structure. Since you didn't post times, I took the liberty to submit your solutions myself, got accepted in 1.499 seconds and 0.782 seconds. While strings are ostensibly immutable, they aren't really, and optimizations in Python/C/OS can keep the string-building with += fast. Commented Feb 26, 2021 at 3:54
• Added a benchmarks answer. Oddly, your two solutions are almost equally fast. Maybe I made some mistake, but I don't see it :-) Commented Feb 26, 2021 at 16:34
• @KellyBundy I'm not entirely surprised by that. My 2nd answer wasn't intended to provide a speed boost. It did seem faster in a "test" I ran (I just multiplied the input text to make it 50000x longer, and used Unix time to measure), but I wasn't rigorous at all.
– FMc
Commented Feb 26, 2021 at 16:41
• Yeah, just changed that answer to say it's probably more odd that the difference was so large on the judge site. Commented Feb 26, 2021 at 16:43
• Yes, I agree ''.join is better. At first, my list of comments didn't even include that even though my code did, I guess I'm so used to it that I forgot I had done it :-). I took another look at your solutions at the judge site. Originally I had submitted them as-is (since they're made for that), now I also submitted the first one after putting the algorithm and all its variables into a function. That made it almost as fast as your second solution there (see comments under my newer answer). So apparently the function-local variables are the main reason your second solution is faster. Commented Feb 26, 2021 at 23:44

Benchmarking a worst-case (key with 1,000 random letters and plaintext with 100,000 random letters).

As @FMc pointed out, isolating the Vigenère algorithm in a function is helpful. Here I do that for benchmarking.

Benchmark results:

74.76 ms  73.98 ms  73.28 ms  original
8.97 ms   8.34 ms   8.01 ms  Kelly_Bundy_1
20.81 ms  20.70 ms  21.24 ms  Kelly_Bundy_2
4.28 ms   4.18 ms   4.10 ms  Kelly_Bundy_3
22.58 ms  22.18 ms  22.62 ms  FMc_1
21.22 ms  20.92 ms  21.62 ms  FMc_2


I'm quite happy to see that my third solution (encrypting entire slices with str.translate) is fastest at this. So apparently the judge site's test suite differs significantly from this worst case.

Benchmark code:

def main():

from timeit import repeat
import random
import string

funcs = original, Kelly_Bundy_1, Kelly_Bundy_2, Kelly_Bundy_3, FMc_1, FMc_2

def random_string(length):
return ''.join(random.choices(string.ascii_uppercase, k=length))

key = random_string(1_000)
plaintext = random_string(100_000)
number = 10

args = key, plaintext

# Correctness
expect = original(*args)
for func in funcs:
result = func(*args)
print(result == expect, func.__name__)
if result != expect:
print(len(expect), len(result), expect[:20], result[:20])
print()

# Speed
tss = [[] for _ in funcs]
for _ in range(3):
for func, ts in zip(funcs, tss):
t = min(repeat(lambda: func(*args), number=number)) / number
ts.append(t)
print(*('%5.2f ms ' % (t * 1e3) for t in ts), func.__name__)
print()

def prep_original():
import string
return string,

def original(key, decrypted_msg, prepped=prep_original()):
string, = prepped

key = (len(decrypted_msg)//len(key) + 1) * key
key = key[0:len(decrypted_msg)]

encrypted_msg = ""

for i in range(len(decrypted_msg)):
key_idx = string.ascii_uppercase.index(key[i])
deciphered_idx = string.ascii_uppercase.index(decrypted_msg[i])

enciphered_idx = (deciphered_idx + key_idx + 1) % len(string.ascii_uppercase)
encrypted_msg += string.ascii_uppercase[enciphered_idx]

return encrypted_msg

def prep_Kelly_Bundy_1():
from itertools import cycle
from string import ascii_uppercase as abc
table = {(abc[i], abc[j]): abc[(i + j + 1) % 26]
for i in range(26)
for j in range(26)}.get
return cycle, table

def Kelly_Bundy_1(key, plaintext, prepped=prep_Kelly_Bundy_1()):
cycle, table = prepped

return ''.join(map(table, zip(plaintext, cycle(key))))

def prep_Kelly_Bundy_2():
from string import ascii_uppercase as abc
from itertools import cycle
index = abc.index
bcd = abc[1:] + abc
return cycle, index, bcd

def Kelly_Bundy_2(key, plaintext, prepped=prep_Kelly_Bundy_2()):
cycle, index, bcd = prepped

key = cycle(map(index, key))
plaintext = map(index, plaintext)

return ''.join([bcd[p + k] for p, k in zip(plaintext, key)])

def prep_Kelly_Bundy_3():
from string import ascii_uppercase as abc
tables = {c: str.maketrans(abc, abc[i:] + abc[:i])
for i, c in enumerate(abc, 1)}
return tables,

def Kelly_Bundy_3(key, plaintext, prepped=prep_Kelly_Bundy_3()):
tables, = prepped

ciphertext = [None] * len(plaintext)
n = len(key)
for i, k in enumerate(key):
ciphertext[i::n] = plaintext[i::n].translate(tables[k])

return ''.join(ciphertext)

def prep_FMc_1():
import string
letters = string.ascii_uppercase
letter_indexes = {c : i for i, c in enumerate(letters)}
n_letters = len(letters)
return letters, letter_indexes, n_letters

def FMc_1(key, decrypted_msg, prepped=prep_FMc_1()):
letters, letter_indexes, n_letters = prepped

key = (len(decrypted_msg)//len(key) + 1) * key
encrypted_chars = []

for i in range(len(decrypted_msg)):
key_idx = letter_indexes[key[i]]
deciphered_idx = letter_indexes[decrypted_msg[i]]

enciphered_idx = (deciphered_idx + key_idx + 1) % n_letters
encrypted_chars.append(letters[enciphered_idx])

enciphered_msg = ''.join(encrypted_chars)
return enciphered_msg

def prep_FMc_2():
import string
from itertools import cycle
letters = string.ascii_uppercase
indexes = {c : i for i, c in enumerate(letters)}
n_letters = len(letters)
return cycle, letters, indexes, n_letters

def FMc_2(key, text, prepped=prep_FMc_2()):
cycle, letters, indexes, n_letters = prepped

encrypted_chars = []
for t, k in zip(text, cycle(key)):
ti = indexes[t]
ki = indexes[k]
ei = (ti + ki + 1) % n_letters
encrypted_chars.append(letters[ei])
return ''.join(encrypted_chars)

main()


Why does this differ from the judge site? Why is my third solution slower than my other two there, when it's faster here? As I explained in my other answer, the intuition for why it could be faster is that the length limits of 1,000 for the key and 100,000 for the plaintext suggest that the plaintext might typically be 100 times longer than the key. Or at least that there are a few test cases at those limits, which would dominate the overall ratio (small cases don't contribute much to the totals and thus to the overall ratio unless there are really many of them). So my third solution is based on the assumption that the total length of all plaintexts divided by the total length of all keys is about 100, making the overhead of using str.translate worth it. But it turns out that the real ratio is only between 2.14 and 2.15, as can be tested with this code that gets accepted (note the assertion at the end):

import sys
from itertools import cycle
from string import ascii_uppercase as abc

table = {(abc[i], abc[j]): abc[(i + j + 1) % 26]
for i in range(26)
for j in range(26)}.get

k = p = 0

lines = map(str.strip, sys.stdin)
for key in lines:
if key == '0':
break
plaintext = next(lines)

k += len(key)
p += len(plaintext)

print(''.join(map(table, zip(plaintext, cycle(key)))))

assert 2.14 < p / k < 2.15

• @Peilonrayz That judge is pretty consistent. I had already noticed that in my previous 43 submissions for this problem, and now I submitted FMc's two solutions again, each five times, alternating between the two (just like I cycle through the solutions in my benchmark, btw). The first solution took 1.372, 1.372, 1.378, 1.439, 1.368 seconds and the second took 0.759, 0.765, 0.775, 0.775, 0.779 seconds. Commented Feb 26, 2021 at 23:23
• Oh wow. Those are some consistent timings. Commented Feb 26, 2021 at 23:25
• @Peilonrayz Submitted FMc's first solution twice more, after modifying it just by putting stuff into a function. Improved to 0.839 and 0.836 seconds that way. Previously I had just submitted it as-is. So apparently most of the speed difference came from the speed of local vs global variables. Commented Feb 26, 2021 at 23:33