算法题

/单词加密。

输入一个字符串和一个非负整数k，对字符串中的每一个字母，
用字母表中其后的第k个字母代替，不够k个时再从字母a循环计数。

例如k=3是，a用d代替，A用D代替，x用a代替，y用b代替，保持大小写不变。

字符串中的非字母字符不变。

字符串的长度不超过100。

输入：一个字符串（无空格）和非负整数k，之间用空格分隔
输出：加密的字符串。

/

# todo 测试输入
str_in = 'asAf$ccz'
str_num = int('3')

# todo 解答
import string


lower_str = string.ascii_lowercase
upper_str = string.ascii_uppercase
str_out, len_str_in = '', len(lower_str)

for i in range(len(str_in)):
    if str_in[i] in lower_str:
        index = lower_str.index(str_in[i]) + str_num

        if index > len_str_in:
            str_out += lower_str[index - len_str_in]
        else:
            str_out += lower_str[index]

    elif str_in[i] in upper_str:
        index = upper_str.index(str_in[i]) + str_num

        if index > len_str_in:
            str_out += upper_str[index - len_str_in]
        else:
            str_out += upper_str[index]

    else:
        str_out += str_in[i]

print(str_out)