3488. 距离最小相等元素查询

src/main.rs

@@ -1,57 +1,47 @@
-use crate::arr::make_matrix;
+use std::collections::HashMap;
 
 mod arr;
 
 struct Solution {}
 impl Solution {
-    pub fn minimum_total_distance(robot: Vec<i32>, factory: Vec<Vec<i32>>) -> i64 {
-        let mut robot = robot.into_iter().map(|x| x as i64).collect::<Vec<i64>>();
-        let mut factory = factory.into_iter().map(|x| x.into_iter().map(|y| y as i64).collect::<Vec<i64>>()).collect::<Vec<Vec<i64>>>();
-        robot.sort_unstable();
-        factory.sort_unstable_by(|a, b| a[0].cmp(&b[0]));
-        factory = factory.into_iter().filter(|f| f[1] > 0).collect();
-        let mut dp = vec![vec![0; robot.len() + 1]; factory.len()];
-        // dp[0][0] = robot[0].abs_diff(factory[0][0]) as i64;
-        for i in 1..=factory[0][1] as usize {
-            dp[0][i] = robot[i - 1].abs_diff(factory[0][0]) as i64 + dp[0][i - 1];
-        }
-        for i in factory[0][1] as usize + 1..=robot.len() {
-            dp[0][i] = 1e11 as i64;
-        }
-        for i in 1..factory.len() {
-            let (f_pos, f_cap) = (factory[i][0], factory[i][1] as usize);
-            for j in 1..=f_cap {
-                dp[i][j] = dp[i][j - 1] + robot[j - 1].abs_diff(f_pos) as i64;
-                dp[i][j] = dp[i][j].min(dp[i - 1][j]);
-                let mut temp = 0;
-                for k in (1..j).rev() {
-                    temp += robot[k].abs_diff(f_pos) as i64;
-                    dp[i][j] = dp[i][j].min(dp[i - 1][k] + temp);
-                }
-            }
-            for j in f_cap + 1..=robot.len() {
-                dp[i][j] = dp[i-1][j];
-                let mut temp = 0;
-                let (mut min_i,mut min_j,mut min_k,mut min_temp,mut min_ans) = (i,j,-1_i64,0,0);
-                for k in (j-f_cap + 1..=j).rev() {
-                    // println!("i: {}, j: {}, k: {}", i, j, k);
-                    temp += robot[k-1].abs_diff(f_pos) as i64;
-                    // println!("temp: {}", temp);
-                    if dp[i - 1][k - 1] + temp < dp[i][j] {
-                        (min_i , min_j,min_k,min_temp,min_ans) = (i,j,k as i64,temp,dp[i - 1][k - 1] + temp)
-                    }
-                    dp[i][j] = dp[i][j].min(dp[i - 1][k - 1] + temp);
-                }
+    pub fn solve_queries(nums: Vec<i32>, queries: Vec<i32>) -> Vec<i32> {
+        let mut mp = HashMap::<i32, Vec<usize>>::new();
+        let n = nums.len() as i32;
+        let mut queries = queries;
+        for i in 0..nums.len() {
+            if mp.contains_key(&nums[i]) {
+                mp.get_mut(&nums[i]).unwrap().push(i);
+            } else {
+                mp.insert(nums[i], vec![i]);
             }
         }
-        dp[factory.len() - 1][robot.len()]
+        for i in 0..queries.len() {
+            let arr = mp.get(&nums[queries[i] as usize]).unwrap();
+            if arr.len() == 1 {
+                queries[i] = -1;
+                continue;
+            }
+
+            let idx = arr.partition_point(|&x| x <= queries[i] as usize);
+            if arr.len() == 2 {
+                queries[i] = (arr[idx % arr.len()] as i32).abs_diff(queries[i]) as i32;
+                queries[i] = queries[i].min(n - queries[i]);
+                continue;
+            }
+            let mut a = arr[idx % arr.len()].abs_diff(queries[i] as usize) as i32;
+            let mut b = arr[(idx + arr.len() - 2) % arr.len()].abs_diff(queries[i] as usize) as i32;
+            a = a.min(n - a);
+            b = b.min(n - b);
+            queries[i] = a.min(b);
+        }
+        queries
     }
 }
 
 fn main() {
     println!(
         "{:?}",
-        Solution::minimum_total_distance(vec![-59,-84], make_matrix("[[-86, 2], [-76, 2], [-46, 2]]"))
+        Solution::solve_queries(vec![12, 19, 12, 8, 12, 10], vec![0, 5, 3, 1])
     );
 }
 /*