【边刷Leetcode边学Rust】537. 复数乘法——实现Trait + 重载操作符

今天，我们通过537. 复数乘法这道题来看看如何在Rust中实现Trait，如何重载操作符。

先来简单分析一下题目。

复数 可以用字符串表示，遵循 "实部+虚部i" 的形式，并满足下述条件：

实部 是一个整数，取值范围是 [-100, 100]

虚部 也是一个整数，取值范围是 [-100, 100]

i² == -1

给你两个字符串表示的复数 num1 和 num2 ，请你遵循复数表示形式，返回表示它们乘积的字符串。

示例：

输入：num1 = "1+-1i", num2 = "1+-1i"

输出："0+-2i"

解释：(1 - i) * (1 - i) = 1 + i² - 2 * i = -2i ，你需要将它转换为 0+-2i 的形式。

有这样几个细节需要注意：

• 即使复数的实部或虚部是0，也不能省略0，如0要写作“0+0i”
• 即使虚部是负数，实部和虚部之间也要用“+”连接，如示例中的“0+-2i”
• 输入和输出都是字符串形式的复数
• 实部和虚部都是范围在 [-100, 100] 内的整数

定义复数类型

首先，我们可以参考Trait std::str::FromStr Examples中Point结构体的定义，定义一个名为Complex的结构体来表示题目中的复数。

struct Complex {
    re: i32,
    im: i32,
}

impl Complex {
    fn new(re: i32, im: i32) -> Self {
        Self { re, im }
    }
}

为了简单起见，这里暂时没有用到泛型，而是直接使用了具体类型i32作为实部和虚部的类型，因为题目限定了实部和虚部的取值类型和范围。

实现FromStr和ToString这两个Trait

既然输入和输出都是字符串，很自然就会想到FromStr和ToString这两个Trait。

姑且可以将Rust中的Trait视作其他语言中的interface。

• FromStr：解析字符串，得到某个类型的值
• ToString：将某个类型的值转换为字符串

先来实现ToString这个较为简单的Trait。

use std::string::ToString;	// ①

// ...

impl ToString for Complex {
    fn to_string(&self) -> String {
        format!("{}+{}i", self.re, self.im)
    }
}

在类型上实现Trait的方法是：在impl关键字之后，提供待实现Trait的名称（这里是ToString，需要先通过use引入该Trait ①），接着是for和类型（这里是Complex）的名称，最后在impl块中，实现定义在Trait中的方法。

接下来用同样的方法实现FromStr：

use std::str::FromStr;	// ①

// ...

#[derive(Debug)]				// ②
struct ParseComplexError;

impl FromStr for Complex {
    type Err = ParseComplexError;		// ③

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let (re, im) = s
            .strip_suffix('i')
            .and_then(|s| s.split_once('+'))
            .ok_or(ParseComplexError)?;

        return Ok(Complex {
            re: re.parse::<i32>().unwrap(),
            im: im.parse::<i32>().unwrap(),
        });
    }
}

这里同样参考了Trait std::str::FromStr Examples中将形如"(1,2)"的坐标转换为Point类型的值的方法。

首先还是通过use引入待实现的Trait FromStr ①。

由FromStr的定义

pub trait FromStr: Sized {
    type Err;

    // Required method
    fn from_str(s: &str) -> Result<Self, Self::Err>;
}

可知，还需要定义一个关联类型（associated type）Err，表示字符串解析失败时的错误。所以我们定义了一个ParseComplexError结构体 ②，并将该结构体作为关联类型 ③。该结构体没有成员，称为基元型结构体（unit-like struct）。

假设题目的输入总是形如“a+bi”的字符串，因此只需先去掉结尾的i（通过strip_suffix('i')），再按+将字符串分为两部分（s.split_once('+')），即可得到复数的实部和虚部（二者都是字符串）。最后再用parse::<i32>()将字符串的实部和虚部转换成整数，就可以构造出Complex结构体的值。

重载“*”操作符

虽然我们也可以通过类似如下的函数来计算两个复数的乘积，

fn multiply(z1: Complex, z2: Complex) -> String {
    let re = z1.re * z2.re - z1.im * z2.im;
    let im = z1.re * z2.im + z1.im * z2.re;
    Complex::new(re, im).to_string()
}

#[test]
fn test_multiply() {
    assert_eq!("0+2i".to_string(), multiply(Complex { re: 1, im: 1 }, Complex { re: 1, im: 1 }));
    assert_eq!("0+-2i".to_string(), multiply(Complex { re: 1, im: -1 }, Complex { re: 1, im: -1 }));
}

但若能直接通过“*”运算符来计算乘积会更加直观，更加优雅。在Rust中，可以通过重载运算符来实现这一点。

use std::ops::Mul;				// ①

// ...

// (a + i b) * (c + i d) == (a*c - b*d) + i (a*d + b*c)
impl Mul<Complex> for Complex {
    type Output = Self;		// ②

    fn mul(self, other: Self) -> Self {
        let re = self.re * other.re - self.im * other.im;
        let im = self.re * other.im + self.im * other.re;
        Self::Output::new(re, im)
    }
}

重载操作符其实还是实现Trait，只要针对Complex类型实现①处导入的Mul Trait，就可以通过形如z1 * z2的表达式来计算复数z1和复数z2的乘积了。

关联类型（associated type）Output是复数乘积的类型，自然也应该是Complex ②。这里的Self指代的就是Complex。

下面我们来测试一下。

fn complex_number_multiply(num1: String, num2: String) -> String {
    let z1 = Complex::from_str(num1.as_str()).unwrap();
    let z2 = Complex::from_str(num2.as_str()).unwrap();
    let z = z1 * z2;
    return z.to_string();
}

#[test]
fn test_complex_number_multiply() {
    assert_eq!("0+2i".to_string(), complex_number_multiply("1+1i".to_string(), "1+1i".to_string()));
    assert_eq!("0+-2i".to_string(), complex_number_multiply("1+-1i".to_string(), "1+-1i".to_string()));
}

可以看到，只需实现这3个Trait，

use std::ops::Mul;
use std::str::FromStr;
use std::string::ToString;

就可以轻松完成“复数乘法”这道题。