RUST GYM Leetcode 1698 Number of Distinct Substrings in a String Rust Solution

1698. Number of Distinct Substrings in a String

Given a string s, return the number of distinct substrings of s.

A substring of a string is obtained by deleting any number of characters (possibly zero) from the front of the string and any number (possibly zero) from the back of the string.

Example 1:

Input: s = "aabbaba"
Output: 21
Explanation: The set of distinct strings is ["a","b","aa","bb","ab","ba","aab","abb","bba","aba","aabb","abba","bbab","baba","aabba","abbab","bbaba","aabbab","abbaba","aabbaba"]

Example 2:

Input: s = "abcdefg"
Output: 28

Constraints:

1 <= s.length <= 500
s consists of lowercase English letters.

Follow up: Can you solve this problem in O(n) time complexity?

Rust Solution

struct Solution;

#[derive(Default)]
struct Trie {
    children: [Option<Box<Trie>>; 26],
}

impl Solution {
    fn count_distinct(s: String) -> i32 {
        let n = s.len();
        let s: Vec<u8> = s.bytes().collect();
        let mut root: Trie = Trie::default();
        let mut res = 0;
        for i in 0..n {
            let mut cur = &mut root;
            for j in i..n {
                if cur.children[(s[j] - b'a') as usize].is_none() {
                    cur.children[(s[j] - b'a') as usize] = Some(Box::new(Trie::default()));
                    res += 1;
                }
                cur = cur.children[(s[j] - b'a') as usize].as_mut().unwrap();
            }
        }
        res
    }
}

#[test]
fn test() {
    let s = "aabbaba".to_string();
    let res = 21;
    assert_eq!(Solution::count_distinct(s), res);
    let s = "abcdefg".to_string();
    let res = 28;
    assert_eq!(Solution::count_distinct(s), res);
}

Having problems with this solution? Click here to submit an issue on github.