rust_algorithms/tree/

trie.rs

//! Trie (prefiks daraxti) — so'zlar lug'ati va avtoto'ldirish uchun.

use std::collections::HashMap;

#[derive(Debug, Default)]
struct Node {
    children: HashMap<char, Node>,
    /// Shu tugunda tugaydigan so'z bormi?
    is_word: bool,
}

/// Trie — har bir qirra bitta harf, ildizdan tugungacha bo'lgan yo'l — prefiks.
///
/// ```text
///        (root)
///        /    \
///      o        k
///      |        |
///      l        i
///     / \       |
///    m   t      t
///    |   |      |
///   [a] [i]    [o]      ← [ ] = so'z shu yerda tugaydi
///        |
///       [n]
///  → "olma", "olti", "oltin", "kito"
/// ```
///
/// | Amal | Murakkablik | Izoh |
/// |---|---|---|
/// | `insert` / `contains` | O(m) | m — so'z uzunligi, **lug'at hajmiga bog'liq emas!** |
/// | `starts_with` | O(m) | prefiks bormi |
/// | `words_with_prefix` | O(m + natija) | avtoto'ldirish |
///
/// **Qayerda ishlatiladi:** qidiruv qatoridagi avtoto'ldirish, telefondagi T9,
/// imlo tekshirgich, IP marshrutlash jadvallari, so'z o'yinlari.
///
/// **HashSet dan farqi:** `HashSet` "bu so'z bormi?" ga javob beradi, Trie esa
/// **prefiks** bo'yicha ham qidira oladi va xotirada umumiy prefikslarni bir marta saqlaydi.
///
/// # Misol
/// ```
/// use rust_algorithms::tree::Trie;
///
/// let mut t = Trie::new();
/// for so_z in ["olma", "olti", "oltin", "kitob"] {
///     t.insert(so_z);
/// }
///
/// assert!(t.contains("olma"));
/// assert!(!t.contains("ol"));        // "ol" — so'z emas, faqat prefiks
/// assert!(t.starts_with("ol"));
///
/// let mut topilgan = t.words_with_prefix("olt");
/// topilgan.sort();
/// assert_eq!(topilgan, vec!["olti", "oltin"]);
/// assert_eq!(t.len(), 4);
/// ```
#[derive(Debug, Default)]
pub struct Trie {
    root: Node,
    len: usize,
}

impl Trie {
    /// Bo'sh trie.
    pub fn new() -> Self {
        Self::default()
    }

    /// So'z qo'shadi. Yangi so'z bo'lsa `true`.
    pub fn insert(&mut self, word: &str) -> bool {
        let mut kursor = &mut self.root;
        for ch in word.chars() {
            kursor = kursor.children.entry(ch).or_default();
        }
        if kursor.is_word {
            return false;
        }
        kursor.is_word = true;
        self.len += 1;
        true
    }

    /// Aynan shu so'z lug'atda bormi?
    pub fn contains(&self, word: &str) -> bool {
        self.node_for(word).is_some_and(|n| n.is_word)
    }

    /// Shu prefiks bilan boshlanadigan so'z bormi?
    pub fn starts_with(&self, prefix: &str) -> bool {
        self.node_for(prefix).is_some()
    }

    /// Prefiks bilan boshlanadigan barcha so'zlar (avtoto'ldirish).
    pub fn words_with_prefix(&self, prefix: &str) -> Vec<String> {
        let mut out = Vec::new();
        if let Some(node) = self.node_for(prefix) {
            let mut buf: Vec<char> = prefix.chars().collect();
            Self::collect(node, &mut buf, &mut out);
        }
        out
    }

    /// Lug'atdagi barcha so'zlar.
    pub fn words(&self) -> Vec<String> {
        self.words_with_prefix("")
    }

    /// So'zni o'chiradi (tugunlar qoladi, faqat belgisi olib tashlanadi).
    pub fn remove(&mut self, word: &str) -> bool {
        let mut kursor = &mut self.root;
        for ch in word.chars() {
            match kursor.children.get_mut(&ch) {
                Some(next) => kursor = next,
                None => return false,
            }
        }
        if !kursor.is_word {
            return false;
        }
        kursor.is_word = false;
        self.len -= 1;
        true
    }

    /// So'zlar soni.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Bo'shmi?
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Berilgan prefiksga mos tugunni topadi.
    fn node_for(&self, prefix: &str) -> Option<&Node> {
        let mut kursor = &self.root;
        for ch in prefix.chars() {
            kursor = kursor.children.get(&ch)?;
        }
        Some(kursor)
    }

    /// Tugundan boshlab barcha so'zlarni yig'adi (DFS).
    fn collect(node: &Node, buf: &mut Vec<char>, out: &mut Vec<String>) {
        if node.is_word {
            out.push(buf.iter().collect());
        }
        for (ch, bola) in node.children.iter() {
            buf.push(*ch);
            Self::collect(bola, buf, out);
            buf.pop();
        }
    }
}

impl<S: AsRef<str>> FromIterator<S> for Trie {
    fn from_iter<I: IntoIterator<Item = S>>(iter: I) -> Self {
        let mut t = Trie::new();
        for s in iter {
            t.insert(s.as_ref());
        }
        t
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn lugat() -> Trie {
        ["olma", "olti", "oltin", "kitob", "kit", "anor"]
            .into_iter()
            .collect()
    }

    #[test]
    fn qoshish_va_qidirish() {
        let t = lugat();
        assert_eq!(t.len(), 6);
        assert!(t.contains("kit"));
        assert!(t.contains("kitob"));
        assert!(!t.contains("kito"));
        assert!(t.starts_with("kito"));
        assert!(!t.starts_with("z"));
    }

    #[test]
    fn takror_qoshilmaydi() {
        let mut t = Trie::new();
        assert!(t.insert("bir"));
        assert!(!t.insert("bir"));
        assert_eq!(t.len(), 1);
    }

    #[test]
    fn prefiks_boyicha_royxat() {
        let t = lugat();
        let mut w = t.words_with_prefix("ol");
        w.sort();
        assert_eq!(w, vec!["olma", "olti", "oltin"]);

        let mut hammasi = t.words();
        hammasi.sort();
        assert_eq!(
            hammasi,
            vec!["anor", "kit", "kitob", "olma", "olti", "oltin"]
        );

        assert!(t.words_with_prefix("zzz").is_empty());
    }

    #[test]
    fn ochirish() {
        let mut t = lugat();
        assert!(t.remove("olti"));
        assert!(!t.contains("olti"));
        assert!(t.contains("oltin")); // uzunroq so'z saqlanib qoldi
        assert!(!t.remove("olti"));
        assert_eq!(t.len(), 5);
    }

    #[test]
    fn bosh_satr_va_unicode() {
        let mut t = Trie::new();
        assert!(t.insert(""));
        assert!(t.contains(""));
        assert!(t.insert("o'zbek"));
        assert!(t.insert("o'zbekiston"));
        assert_eq!(t.words_with_prefix("o'zbeki"), vec!["o'zbekiston"]);
    }
}