3
\$\begingroup\$

I just started learning Rust, and this is my solution to a problem from Atcoder's Regular Contest #067, D - Walk and Teleport.


The problem can be summarized as follows:

  • There are N towns in a line, and you are about to visit all of them.
  • Town i is located at the point with coordinate xi.
  • You can travel either by walking or by teleporting. (You can combine them)
    • Walking costs you a * |Δx|, where a is a constant given and Δx is the distance you are about to travel.
    • Teleporing to any location costs you b, regardless of the distance you travel.
  • Given a, b and a sorted list of coordinates (xi), find the minimum cost to travel all towns.

Input is given in the form of

N a b
x1 x2 ... xN

where N is the number of cities and a, b, xi is what described above. Again, xi is sorted. All a, b, xi are integers in the range of [1, 109].


My speculation is that you can just travel from left to right, and choose the cheaper way to travel to the next town (a*(xj+1-xj) versus b).

This is my code.

use std::io;
use std::cmp::min;

fn main (){
    let input_one = read_ints();
    let a = input_one[1]; let b = input_one[2];
    let xs = read_ints();
    println!("{}", solve(&a,&b,&xs));
}

fn distances(xs:&[u64]) -> Vec<u64>{
    // Find the distances between adjacent towns
    xs.iter().zip(xs.iter().skip(1)).map(|(a,b)| b-a).collect()
}

fn solve(a:&u64, b:&u64, xs:&[u64]) -> u64 {
    let dists = distances(xs);
    dists.iter().fold(0, |acc,dist| {
        acc + min(a*dist, *b)
    })
}
fn read_ints() -> Vec<u64> {
    get_line().trim().split_whitespace().map(|n| n.parse().unwrap()).collect()
}

fn get_line() -> String {
    let mut buf = String::new();
    io::stdin().read_line(&mut buf).ok();
    buf
}

As I'm new to Rust, my main concern is the idiomaticness of the code. I've been writing Haskell and Python and there are many "it works but I'm not sure if I'm doing it right" moments. (Like, &[u64]'s were &Vec<u64> just until I was about to post this question.)

Any comments, including styles, naming, "loop is preferred", or anything are very much appreciated. Thank you.

\$\endgroup\$
2
\$\begingroup\$
  1. Learn to love rustfmt.

    1. No space between a method name and the parens, space after the parens.

      -fn main (){
      +fn main() {
      
    2. Multiple statements belong on different lines.

      -    let a = input_one[1]; let b = input_one[2];
      +    let a = input_one[1];
      +    let b = input_one[2];
      
    3. Spaces go between comma-separated items.

      -    println!("{}", solve(&a,&b,&xs));
      +    println!("{}", solve(&a, &b, &xs));
      
    4. Spaces go after the : in declarations.

      -fn distances(xs:&[u64]) -> Vec<u64>{
      +fn distances(xs: &[u64]) -> Vec<u64> {
      
  2. Using Result::ok is not appropriate error handling. If there had been a failure while reading from standard input, your program would continue with an empty string. That's unlikely to be what you want.

  3. unwrapping a Result is better than losing an error, expecting is better than that as you get a little bit of a clue as to where there was a problem, and returning a Result is even better.

  4. If you are going to document a function, use doc comments, not programmer comments.

  5. Zipping the slice with itself is fine, but using slice::windows is a bit shorter, and may be easier to understand when reading the code.

  6. folding over the distances is fine, but a combination of map and sum is probably easier to understand and is as efficient.

  7. There's no real benefit to taking values like u64 by reference. They are small and implement Copy. I'd expect that passing by value would be the same speed if not faster.

  8. a, b, 'xs' are bad variable names. Make them something useful.

use std::io;
use std::cmp::min;

fn main() {
    let input_one = read_ints();
    let walk_cost = input_one[1];
    let teleport_cost = input_one[2];
    let locations = read_ints();
    println!("{}", solve(walk_cost, teleport_cost, &locations));
}

/// Find the distances between adjacent towns
fn distances(locations: &[u64]) -> Vec<u64> {
    locations.windows(2).map(|w| w[1] - w[0]).collect()
}

fn solve(walk_cost: u64, teleport_cost: u64, locations: &[u64]) -> u64 {
    let dists = distances(locations);
    dists.iter().map(|dist| min(walk_cost * dist, teleport_cost)).sum()
}

fn read_ints() -> Vec<u64> {
    get_line()
        .trim()
        .split_whitespace()
        .map(|n| n.parse().expect("Unable to parse"))
        .collect()
}

fn get_line() -> String {
    let mut buf = String::new();
    io::stdin().read_line(&mut buf).expect("Unable to read");
    buf
}

Overall, your code was easy to read and understand.


However, it feels like there's more allocation than there needs to be. Specifically, taking an iterator and making a Vec to make it into an iterator to make a Vec to just iterate through it again seems unfortunate.

One silly alternative would be something like this:

extern crate itertools;

use std::io::{self, BufRead};
use std::cmp::min;

use itertools::Itertools;

fn main() {
    let input = io::stdin();
    let input = input.lock();
    let mut lines = input.lines();

    let line = lines.next().expect("no line").expect("Unable to read line");
    let mut config_line = line.split_whitespace().map(|n| n.parse::<u64>().expect("not an integer")).skip(1);
    let walk_cost = config_line.next().expect("no walk cost provided");
    let teleport_cost = config_line.next().expect("no teleport cost provided");

    let line = lines.next().expect("no line").expect("Unable to read line");
    let locations = line.split_whitespace().map(|n| n.parse::<u64>().expect("not an integer"));
    let distances = locations.tuple_windows().map(|(a, b)| b - a);
    let costs = distances.map(|dist| min(walk_cost * dist, teleport_cost));
    let total_cost: u64 = costs.sum();

    println!("{}", total_cost);
}

This is not what I'd consider idiomatic Rust code, however. I always find iterators that return Results to be a bit awkward to deal with while still passing up errors to the calling function.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.