🚀 Отслеживание цветов шариков с использованием 2 HashMap

leetcoder, изменено 09-02-2025 ,

244

Ссылка на задачу – 3160. Find the Number of Distinct Colors Among the Balls.

Описание задачи 😊

Даны шарики с уникальными метками от 0 до limit (всего limit + 1 шариков)
И список запросов, где каждый запрос имеет вид [x, y]
При каждом запросе шарик с меткой x окрашивается в цвет y.
После каждого запроса необходимо вернуть количество различных цветов, присутствующих среди раскрашенных шариков (неокрашенные не учитываются).

Идея 💡

Основная идея заключается в использовании двух хеш-таблиц:

ball_to_color – для хранения текущего цвета каждого шарика.
color_counts – для подсчёта количества шариков каждого цвета.

Это позволяет за константное (в среднем) время обновлять информацию при каждом запросе и быстро определять число уникальных цветов.

Детали подхода 🔍

Обновление цвета шарика:
- При выполнении запроса, если шарик уже был раскрашен, получаем его старый цвет и уменьшаем счётчик в color_counts. Если счётчик для старого цвета становится равным нулю, уменьшаем количество уникальных цветов.
Установка нового цвета:
- Обновляем запись в ball_to_color новым цветом. Затем увеличиваем счётчик для нового цвета в color_counts. Если новый цвет появляется впервые (счётчик становится равным 1), увеличиваем количество уникальных цветов.
Формирование результата:
- После каждого запроса записываем текущее число уникальных цветов в результирующий вектор.

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

Время работы:
O(Q), где Q — количество запросов. Каждая операция над HashMap выполняется в среднем за O(1).
Память:
O(N), где N — количество окрашиваемых шариков или уникальных цветов в худшем случае.

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

use std::collections::HashMap;
use std::collections::hash_map::Entry;

impl Solution {
    pub fn query_results(limit: i32, queries: Vec<Vec<i32>>) -> Vec<i32> {
        let mut ball_to_color = HashMap::new();
        let mut color_counts = HashMap::new();
        let mut distinct_colors = 0;
        let mut result = Vec::with_capacity(queries.len());

        for query in queries {
            let (ball, new_color) = (query[0], query[1]);

            // If the ball was already painted, update the count for the old color.
            if let Some(old_color) = ball_to_color.insert(ball, new_color) {
                let Entry::Occupied(count) = color_counts.entry(old_color)
                    .and_modify(|count| *count -= 1) else {
                    unreachable!("color should be accounted")
                };
                if *count.get() == 0 {
                    distinct_colors -= 1;
                }
            }

            // Update the count for the new color.
            if 1 == *color_counts.entry(new_color)
                .and_modify(|count| *count += 1)
                .or_insert(1) 
            {
                distinct_colors += 1;
            }

            result.push(distinct_colors);
        }

        result
    }
}

Tags: #rust #algorithms #hashmap