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 = sys.stdin.readline().strip()
    key = int(key) if key.isdigit() else key
    if key == 0: break

    decrypted_msg = sys.stdin.readline().strip()
    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 = &quot;&quot;

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

decrypted_msg is misleading. It's not decrypted but rather &quot;not encrypted&quot;. But really just call it &quot;plaintext&quot;, as the problem statement does.
Similarly, &quot;enciphered&quot; and &quot;encrypted&quot; are not just inconsistent (make up your mind), but go against the problem statement's term &quot;ciphertext&quot;.
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 &quot;B&quot;, 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')