diff --git a/src/main.rs b/src/main.rs index 4182ca5..5ecada7 100644 --- a/src/main.rs +++ b/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, factory: Vec>) -> i64 { - let mut robot = robot.into_iter().map(|x| x as i64).collect::>(); - let mut factory = factory.into_iter().map(|x| x.into_iter().map(|y| y as i64).collect::>()).collect::>>(); - 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, queries: Vec) -> Vec { + let mut mp = HashMap::>::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]) ); } /* @@ -67,4 +57,4 @@ fn main() { -86 -74 -46 -29 -84 -60 -59 -7 -5 -3 - */ \ No newline at end of file + */