3318. 计算子数组的 x-sum I

src/main.rs

@@ -1,71 +1,106 @@
+use std::cmp::min;
+use std::collections::{BTreeSet, HashMap};
+
 mod arr;
-struct Bank {
-    _balance: Vec<i64>,
+
+#[derive(Debug)]
+struct Pair {
+    index: i32,
+    value: i32,
 }
-
-
-/**
- * `&self` means the method takes an immutable reference.
- * If you need a mutable reference, change it to `&mut self` instead.
- */
-impl Bank {
-
-    fn new(balance: Vec<i64>) -> Self {
-        Bank { _balance: balance }
+impl PartialEq for Pair {
+    fn eq(&self, other: &Pair) -> bool {
+        self.value == other.value && self.index == other.index
     }
-
-    fn valid(&self, account: i32,money: i64) -> bool {
-        account > 0 && account as usize <= self._balance.len() && self._balance[account as usize - 1] >= money
+}
+impl Eq for Pair {}
+impl PartialOrd for Pair {
+    fn partial_cmp(&self, other: &Pair) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
     }
-    fn transfer(&mut self, account1: i32, account2: i32, money: i64) -> bool {
-        if !self.valid(account1, money) {
-            return false;
-        }
-        self._balance[account1 as usize - 1] -= money;
-        self._balance[account2 as usize - 1] += money;
-        true
-    }
-
-    fn deposit(&mut self, account: i32, money: i64) -> bool {
-        if !self.valid(account, 0) {
-            return false;
-        }
-        self._balance[account as usize - 1] += money;
-        true
-    }
-
-    fn withdraw(&mut self, account: i32, money: i64) -> bool {
-        if !self.valid(account, money) {
-            false
+}
+impl Ord for Pair {
+    fn cmp(&self, other: &Pair) -> std::cmp::Ordering {
+        if self.value == other.value {
+            other.index.cmp(&self.index)
         } else {
-            self._balance[account as usize - 1] -= money;
-            true
+            other.value.cmp(&self.value)
         }
     }
 }
 
-/**
- * Your Bank object will be instantiated and called as such:
- * let obj = Bank::new(balance);
- * let ret_1: bool = obj.transfer(account1, account2, money);
- * let ret_2: bool = obj.deposit(account, money);
- * let ret_3: bool = obj.withdraw(account, money);
- */
+impl Clone for Pair {
+    fn clone(&self) -> Self {
+        Pair {
+            index: self.index,
+            value: self.value,
+        }
+    }
+}
+
+struct Solution;
+impl Solution {
+    pub fn find_x_sum(nums: Vec<i32>, k: i32, x: i32) -> Vec<i32> {
+        let x = x as usize;
+        let mut tree: BTreeSet<Pair> = BTreeSet::new();
+        let mut dict: HashMap<i32, Pair> = HashMap::new();
+        let mut ans: Vec<i32> = vec![];
+        for i in 0..k as usize {
+            let mut pair = Pair {
+                index: nums[i],
+                value: 1,
+            };
+            if dict.contains_key(&nums[i]) {
+                let exist_pair = dict.get(&nums[i]).unwrap();
+                tree.remove(exist_pair);
+                pair.value += exist_pair.value;
+            }
+            dict.insert(nums[i], pair.clone());
+            tree.insert(pair);
+        }
+        ans.push(
+            tree.iter()
+                .take(min(tree.len(), x))
+                .fold(0, |last, cur| last + cur.value * cur.index),
+        );
+        for i in 0..nums.len() - k as usize {
+            let mut left_item = dict[&nums[i]].clone();
+            tree.remove(&left_item);
+            if left_item.value > 1 {
+                left_item.value -= 1;
+                dict.insert(nums[i], left_item.clone());
+                tree.insert(left_item);
+            } else {
+                dict.remove(&nums[i]);
+            }
+            if dict.contains_key(&nums[i + k as usize]) {
+                let mut pair = dict.get(&nums[i + k as usize]).unwrap().clone();
+                tree.remove(&pair);
+                pair.value += 1;
+                dict.insert(nums[i + k as usize], pair.clone());
+                tree.insert(pair);
+            } else {
+                let pair = Pair {
+                    index: nums[i + k as usize],
+                    value: 1,
+                };
+                dict.insert(nums[i + k as usize], pair.clone());
+                tree.insert(pair);
+            };
+            ans.push(
+                tree.iter()
+                    .take(min(tree.len(), x))
+                    .fold(0, |last, cur| last + cur.value * cur.index),
+            );
+        }
+        ans
+    }
+}
 
 fn main() {
-    let oper = ["withdraw", "transfer", "deposit", "transfer", "withdraw"];
-    let val = arr::make_matrix("[[3, 10], [5, 1, 20], [5, 20], [3, 4, 15], [10, 50]]");
-    //输出：
-    //[true, true, true, false, false]
-    let mut bank = Bank::new(vec![10, 100, 20, 50, 30]);
-    for i in 0..oper.len() {
-        let result = match oper[i] {
-            "withdraw" => bank.withdraw(val[i][0], val[i][1] as i64),
-            "transfer" => bank.transfer(val[i][0], val[i][1], val[i][2] as i64),
-            "deposit" => bank.deposit(val[i][0], val[i][1] as i64),
-            _ => false,
-        };
-        println!("{}", result);
-    }
-
+    let nums = vec![1, 1, 2, 2, 3, 4, 2, 3];
+    let k = 6;
+    let x = 2;
+    let result = Solution::find_x_sum(nums, k, x);
+    println!("{:?}", result);
 }