leetcode

2463-minimum-total-distance-traveled.rs (5397B)

/// 2463. Minimum Total Distance Traveled (Hard)
///
/// There are some robots and factories on the X-axis. You are given an integer
/// array robot where robot[i] is the position of the ith robot. You are also
/// given a 2D integer array factory where factory[j] = [positionj, limitj]
/// indicates that positionj is the position of the jth factory and that the jth
/// factory can repair at most limitj robots.
///
/// The positions of each robot are *unique*. The positions of each factory are
/// also *unique*. Note that a robot can be *in the same position* as a factory
/// initially.
///
/// All the robots are initially broken; they keep moving in one direction. The
/// direction could be the negative or the positive direction of the X-axis.
/// When a robot reaches a factory that did not reach its limit, the factory
/// repairs the robot, and it stops moving.
///
/// *At any moment*, you can set the initial direction of moving for *some*
/// robot. Your target is to minimize the total distance traveled by all the
/// robots.
///
/// Return the minimum total distance traveled by all the robots. The test cases
/// are generated such that all the robots can be repaired.
///
/// *Note that*
///
///   * All robots move at the same speed.
///   * If two robots move in the same direction, they will never collide.
///   * If two robots move in opposite directions and they meet at some point,
///   they do not collide. They cross each other.
///   * If a robot passes by a factory that reached its limits, it crosses it as
///   if it does not exist.
///   * If the robot moved from a position x to a position y, the distance it
///   moved is |y - x|.
///
/// *Example 1:*
///
///   *Input:* robot = [0,4,6], factory = [[2,2],[6,2]]
///   *Output:* 4
///   *Explanation:* As shown in the figure:
///   - The first robot at position 0 moves in the positive direction. It will be
///     repaired at the first factory.
///   - The second robot at position 4 moves in the negative direction. It will be
///     repaired at the first factory.
///   - The third robot at position 6 will be repaired at the second factory. It
///     does not need to move.
///   The limit of the first factory is 2, and it fixed 2 robots.
///   The limit of the second factory is 2, and it fixed 1 robot.
///   The total distance is |2 - 0| + |2 - 4| + |6 - 6| = 4. It can be shown that
///   we cannot achieve a better total distance than 4.
///
/// *Example 2:*
///
///   *Input:* robot = [1,-1], factory = [[-2,1],[2,1]]
///   *Output:* 2
///   *Explanation:* As shown in the figure:
///   - The first robot at position 1 moves in the positive direction. It will be
///     repaired at the second factory.
///   - The second robot at position -1 moves in the negative direction. It will be
///     repaired at the first factory.
///   The limit of the first factory is 1, and it fixed 1 robot.
///   The limit of the second factory is 1, and it fixed 1 robot.
///   The total distance is |2 - 1| + |(-2) - (-1)| = 2. It can be shown that we
///   cannot achieve a better total distance than 2.
///
/// *Constraints:*
///
///   * 1 <= robot.length, factory.length <= 100
///   * factory[j].length == 2
///   * -109 <= robot[i], positionj <= 109
///   * 0 <= limitj <= robot.length
///   * The input will be generated such that it is always possible to repair
///   every robot.
///
use leetcode::*;

struct Solution {}

impl Solution {
    fn dp(
        robot: &Vec<i32>,
        factory: &Vec<Vec<i32>>,
        cache: &mut [i64],
        ri: usize,
        fi: usize,
    ) -> i64 {
        if ri >= robot.len() {
            return 0;
        }

        if fi >= factory.len() {
            return i64::MAX;
        }

        let cached = cache[ri * factory.len() + fi];
        if cached != -1 {
            return cached;
        }
        let mut answer = Self::dp(robot, factory, cache, ri, fi + 1);
        let mut dist = 0;
        for i in 0..factory[fi][1] {
            let rri = ri + i as usize;
            if rri >= robot.len() {
                break;
            }
            dist += (robot[rri] - factory[fi][0]).abs() as i64;
            // we skip the current factory entirely. it is never optimal for
            // non-consecutive robots to be assigned to the same factory (!)
            let res = Self::dp(robot, factory, cache, rri + 1, fi + 1);
            if res == i64::MAX {
                continue;
            }
            answer = std::cmp::min(answer, dist + res)
        }
        cache[ri * factory.len() + fi] = answer;

        answer
    }

    pub fn minimum_total_distance(mut robot: Vec<i32>, mut factory: Vec<Vec<i32>>) -> i64 {
        factory.sort_by(|a, b| a[0].cmp(&b[0]));
        robot.sort();

        let mut cache: Box<[i64]> = vec![-1; factory.len() * robot.len()].into_boxed_slice();
        Self::dp(&robot, &factory, &mut cache, 0, 0)
    }
}

pub fn main() {
    println!(
        "{:?}",
        Solution::minimum_total_distance(arg_into(1), arg_into(2))
    )
}

#[cfg(test)]
mod tests {
    use crate::Solution;
    use leetcode::{vi, vvi};

    #[test]
    fn test() {
        assert_eq!(
            Solution::minimum_total_distance(vi("[0,4,6]"), vvi("[[2,2],[6,2]]")),
            4
        );
        assert_eq!(
            Solution::minimum_total_distance(vi("[1,-1]"), vvi("[[-2,1],[2,1]]")),
            2
        );
    }
}