🚀 Подсчёт хороших троек с помощью отображения индексов и дерева Фенвика

leetcoder, 15-04-2025 ,

Ссылка на задачу с LeetCode — 2179. Count Good Triplets in an Array.

📘 Условие задачи

Даны два массива nums1 и nums2, являющиеся перестановками чисел от 0 до n-1.
Нужно найти количество хороших троек (x, y, z), таких что:

индексы этих значений возрастают в обоих массивах:
pos₁(x)<pos₁(y)<pos₁(z) и pos₂(x)<pos₂(y)<pos₂(z),
где pos₁(v) и pos₂(v) — позиция значения v в nums1 и nums2 соответственно.

📌 Пример

Вход: nums1 = [2,0,1,3], nums2 = [0,1,2,3]
Выход: 1
Пояснение:
- Существуют 4 тройки, удовлетворяющие порядку в nums1:
  (2,0,1), (2,0,3), (2,1,3), (0,1,3)
- Из них только (0,1,3) также удовлетворяет порядку в nums2.

💡 Идея

Рассматриваем каждое число как средний элемент потенциальной хорошей тройки.
Если знать, сколько подходящих чисел было до него и после него, можно посчитать количество троек с этим числом посередине.

Дерево Фенвика и есть подходящая структура данных, позволяющая динамически поддерживать префиксные суммы, за счёт которых мы эффективно рассчитываем необходимые величины.

🛠 Подробности подхода

Создаём отображение: значение → индекс в nums2.
Идём по nums1. Пусть y — текущий элемент, а pos₁(y) и pos₂(y) — его позиции в обоих массивах.
- Считаем, сколько уже встреченных значений были до него в обоих массивах — это left_count.
- Считаем, сколько значений ещё впереди, которые также находятся после него и в nums2:
  right_count = (n−1−pos₂(y)) − (pos₁(y)−left_count)</code
- Добавляем left_count * right_count к результату.
Обновляем дерево Фенвика, чтобы отметить текущую позицию y как «увиденную».

⏱ Асимптотики

Время: O(n·log n)
- Построение отображения: O(n)
- Для каждого элемента — query и add в дереве Фенвика: O(log n)
Память: O(n)
- Используется для массива отображения и дерева Фенвика

🧾 Исходный код

struct FenwickTree {
    tree: Vec<i32>,
}

impl Solution {
    pub fn good_triplets(nums1: Vec<i32>, nums2: Vec<i32>) -> i64 {
        let n = nums1.len();

        // Map value → index in nums2
        let mut pos2 = vec![0; n];
        for (i, v) in nums2.into_iter().enumerate() {
            pos2[v as usize] = i;
        }

        let mut sum_tree = FenwickTree::new(n);

        nums1.into_iter().enumerate().map(|(i, v)| {
            let j = pos2[v as usize]; // Corresponding index in nums2
            // Count values before i in nums1 that are also before j in nums2
            let left_count = sum_tree.query(j) as usize;
            // Count values after i in nums1 that are also after j in nums2
            let right_count = n - j - 1 - (i - left_count);
            // Mark position j as seen
            sum_tree.add(j, 1);

            (left_count * right_count) as i64
        }).sum()
    }
}

impl FenwickTree {
    // Initializes a Fenwick Tree with the given size.
    fn new(size: usize) -> Self {
        let n = 1 << (1 + Self::msb(size));
        Self { tree: vec![0; n + 1] }
    }

    // Adds a value to the Fenwick Tree at a specific index.
    fn add(&mut self, mut index: usize, value: i32) {
        index += 1;
        while index < self.tree.len() {
            self.tree[index] += value;
            index += index & (!index + 1);
        }
    }

    // Queries the cumulative sum up to a specific index.
    fn query(&self, mut index: usize) -> i32 {
        let mut sum = 0;
        index += 1;
        while index > 0 {
            sum += self.tree[index];
            index &= index - 1;
        }
        sum
    }

    fn msb(mut n: usize) -> usize {
        let mut pos: usize = 0;
        while n > 0 {
            n >>= 1;
            pos += 1;
        }
        pos.saturating_sub(1)
    }
}

Tags: #rust #algorithms #fenwick