用Python定义一个函数,实现对一个复数a+bj求模的平方并返回该值

紫光的一个笔试题目，求复数模的平方。和同学交流了好一会后写出如下结果

a = input("input")
b =a.split("+" or "-")  #按照符号分割后赋值给b
def xs(b):
    if len(b)>1:        #实部与虚部都存在
        if "j" in b[0] :
            s = int(b[0][:-1])
            x= int(b[1])
        elif "j" in b[1]:
            x = int(b[1][:-1])
            s = int(b[0])
        return s * s + x * x
    elif "j" in b[0]:   #只存在虚部
        s = int(b[0][:-1])
        return s*s
    else:               #只存在实部
        return pow(int(b[0]),2)
print(xs(b))

复数模的数学模型是对实部与虚部求平方后相加，再开根号如：a = 1-2j

（1*1+（-2）*（-2））**0.5

 

elif "j" in b[0]:   #只存在虚部
    s = int(b[0][:-1])

这一行中使用b[0]的原因，如果只存在实部或虚部，使用s = int(b[:-1])切割["200j"]的时候获取为null.所以即使只存在实部或虚部也需要使用b[0]