LeetCodeでRustを学ぶ【Two Sum】

問題

https://leetcode.com/problems/two-sum/

問題文

Given an array of integers nums and an integer target, return indices of the two numbers such that they add up to target.
You may assume that each input would have exactly one solution, and you may not use the same element twice.
You can return the answer in any order.

制約

$2 <= nums.length <= 10^4$
$-10^9 <= nums[i] <= 10^9$
$-10^9 <= target <= 10^9$
Only one valid answer exists.

実装例

$O(N^2)$の解法

impl Solution {
    pub fn two_sum(nums: Vec<i32>, target: i32) -> Vec<i32> {
        let mut ans = Vec::with_capacity(2);
        for (idx1, num1) in nums.iter().enumerate() {
            for (idx2, num2) in nums.iter().enumerate() {
                if idx1 == idx2 {
                    continue;
                }
                if num1 + num2 == target {
                    ans.push(idx1 as i32);
                    ans.push(idx2 as i32);
                    return ans;
                }
            }
        }
        ans
    }
}

$O(N)$の解法

use std::collections::HashMap;

impl Solution {
    pub fn two_sum(nums: Vec<i32>, target: i32) -> Vec<i32> {
        let mut map = HashMap::new();
        let mut ans = Vec::with_capacity(2);

        for (idx, num) in nums.into_iter().enumerate() {
            let check = map.get(&num);
            match check {
                Some(&val) => {
                    ans.push(val);
                    ans.push(idx as i32);
                    break;
                },
                None => _ = map.insert(target-num, idx as i32),
            }
        }
        ans
    }
}

Rustワンポイント

ベクタの要素数が事前にわかっている場合は，Vec::with_capacity() を使って容量を先に確保しよう！
今回の例ではほとんど差はないけれど，パフォーマンスが良くなるよ！
iter()を使うか，into_iter()を使うかは所有権の移動と参照渡しと値渡しを考えて使い分けよう！