Cleanly modifying values in hashmap for a k-sum challenge

Question

Just started learning Rust and here's a program that has caused me some agony to implement in Rust. This is very easy to implement in C++, but Rust's borrow checker is causing all sorts of trouble. My program tries to solve the following challenge:

You are given an integer array nums and an integer k.

In one operation, you can pick two numbers from the array whose sum equals k and remove them from the array.

Return the maximum number of operations you can perform on the array.

The algorithm I'm using is the obvious one:

1. Scan over the input making a hash map of the counts of each integer in the input.
2. Loop over the integers and for each n check if the count of n and k - n is more than 0 (more than 1 if n and k-n are equal).
3. If yes, decrease counts of both integers and increment the result.

My implementation in Rust is the following, which seems quite ugly:

use std::collections::HashMap;
use std::collections::hash_map::Entry;

impl Solution {
    pub fn max_operations(nums: Vec<i32>, k: i32) -> i32 {
        let mut result = 0;
        
        // Keep track of counts of each integer.
        let mut counts = HashMap::<i32,i32>::new();
        
        for n in nums.iter() {
            let count = counts.entry(*n).or_insert(0);
            *count += 1;
        }
        
        // Loop over each element and check if it has a matching pair to remove.
        for n in nums.iter() {
            if *n == k - *n {
                if *counts.entry(*n).or_insert(0) < 2 {
                    continue;
                }
                result += 1;
                *counts.entry(*n).or_insert(0) -= 2;
                continue;
            }
            
            match counts.entry(*n) {
                Entry::Occupied(o) => {
                    if *o.get() < 1 {
                        continue;
                    }
                    match counts.entry(k - *n) {
                        Entry::Occupied(o) => {
                            if *o.get() < 1 {
                                continue;
                            }
                        },
                        Entry::Vacant(_) => continue    
                    }
                },
                Entry::Vacant(_) => continue
                
            }
                             
            result += 1;
            *counts.entry(*n).or_insert(0) -= 1;
            *counts.entry(k - *n).or_insert(0) -= 1;   
        }
        return result;
    }
}

What I would like to do is something like the following instead of the match block and everything that comes after it:

if let counts.get(*n) > 0 && let counts.get(k - *n) > 0 {
    *counts.entry(*n).or_insert(0) -= 1;
    *counts.entry(k - *n).or_insert(0) -= 1;   
}

This is not possible and it would also cause unnecessary zeroes to be inserted in the hashmap for elements not present.

Does Rust have any way of writing code like this in a more compact way?

Answer 1

Destructuring in for loop

Both of your for loops are defined as

for n in nums.iter() { .. }

This makes n have type &i32 so you have to dereference it with *n every time you attempt to use it. Since i32 implements Copy, you could instead use

for &n in nums.iter() {
    ..
}

This destructures the reference so n itself has type i32, making it much more ergonomic to use.

Unnecessary Entry API

The HashMap Entry API is very useful when you want to modify the map and don't know much about the keys. A perfect example is how you initialize the map with or_insert.

However, in the rest of your code it would be much more ergonomic to instead rely on get and Index which return a reference to the value in a HashMap given the key, e.g.

counts.get(key) // Returns Option<&i32> with Some(value) corresponding to the key if it exists, otherwise None.
counts[key]     // Returns &i32 with the value, PANICS if the key does not exist.

Since you know keys exist for all nums, you can simply use the indexing approach. For k - n, use counts.get() instead. Instead of the match statement, the following will work

if let Some(&cnt) = counts.get(&(k - n)) {
    if counts[&n] > 0 && cnt > 0 {
        result += 1;
        *counts.get_mut(&n).unwrap() -= 1;
        *counts.get_mut(&(k - n)).unwrap() -= 1;
    }
}

Note that indexing doesn't work if you need to mutate the map since HashMap does not implement IndexMut. You will need to use get_mut instead, but you can unwrap the result since we know the key exists.

We cannot statically prove to the compiler that get_mut(n) and get_mut(k - n) will correspond to different keys so we need to do a bit of a dance to avoid overlapping mutable references.

Final code

fn max_operations(nums: Vec<i32>, k: i32) -> i32 {
    let mut counts = HashMap::<i32, i32>::new();
    for &n in nums.iter() {
        *counts.entry(n).or_insert(0) += 1;
    }

    let mut result = 0;
    for &n in nums.iter() {
        if n == k - n {
            if counts[&n] >= 2 {
                result += 1;
                *counts.get_mut(&n).unwrap() -= 2;
            }
        } else if let Some(&cnt) = counts.get(&(k - n)) {
            if counts[&n] > 0 && cnt > 0 {
                result += 1;
                *counts.get_mut(&n).unwrap() -= 1;
                *counts.get_mut(&(k - n)).unwrap() -= 1;
            }
        }
    }

    result
}

Answer 2

You have a great implementation-level answer by apilat. I'll address the algorithm:

• Scan over the input making a hash map of the counts of each integer in the input.

• Loop over the integers and for each n check if the count of n and k - n is more than 0 (more than 1 if n and k-n are equal).

• If yes, decrease counts of both integers and increment the result.

The question asks how many pairs can be removed from the array, not how many distinct pairs. So we can increase the count by min(count(n), count(k-n)). And we don't need to modify the counts if we're not going to look at them again. Rust isn't my language, but I would expect to be able to loop over the entries, ignoring those where n > k - n, and special-casing n == k - n to contribute count(n) / 2 to the sum.

BTW, a note on terminology: k-sum generally means "select k numbers whose sum is a given value". So this is one kind of 2-sum problem.