3583. 统计特殊三元组

src/main.rs

@@ -1,19 +1,30 @@
+use std::collections::HashMap;
+
 struct Solution;
 mod arr;
 impl Solution {
-    pub fn count_permutations(complexity: Vec<i32>) -> i32 {
+    pub fn special_triplets(nums: Vec<i32>) -> i32 {
+        let mut map: HashMap<i32, i32> = HashMap::new();
+        let mut map2: HashMap<i32, i32> = HashMap::new();
         let md = 1e9 as i64 + 7;
-        let mut ans = 1_i64;
-        for i in 1..complexity.len() {
-            if complexity[i] <= complexity[0] {
-                return 0;
+        let mut ans = 0_i64;
+        nums.iter().for_each(|x| {
+            map.insert(*x, map.get(&x).copied().unwrap_or(0) + 1);
+        });
+        map2.insert(nums[0], 1);
+        for i in 1..nums.len() - 1 {
+            let mp2d = map2.get(&(nums[i] * 2)).copied().unwrap_or(0) as i64;
+            let mut mp1d = map.get(&(nums[i] * 2)).copied().unwrap_or(0) as i64 - mp2d;
+            if nums[i] == nums[i] * 2 {
+                mp1d -= 1;
             }
-            ans = (ans * i as i64) % md;
+            ans = (ans + mp1d * mp2d) % md;
+            map2.insert(nums[i], map2.get(&nums[i]).copied().unwrap_or(0) + 1);
         }
         ans as i32
     }
 }
 fn main() {
-    let result = Solution::count_permutations(vec![1, 3, 3, 4, 4, 4]);
+    let result = Solution::special_triplets(vec![0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]);
     println!("{:?}", result);
 }