#StackBounty: #beginner #algorithm #rust Memory/Time usage on substring search code

Bounty: 50

I have implemented a substring search, which works as expected and provides the right output. On some edge cases tho, I get either Memory or Time limit exceeded because:

  1. the string length is too big
  2. there are too many queries

The code asks for a few inputs:

  1. a string
  2. the number of queries
  3. query

Info about the string and the queries:

  1. string length is not bigger than 300.000 and consists of English lowercase characters.
  2. max number of queries is 300.000
  3. sum of the queries length is not bigger than 300.000

I get MLE if I use HashMap caching, and I get TLE without it.

Where can my code be improved, so it also handles those edge cases?

use std::collections::HashMap;
use std::io::{stdin, stdout, BufWriter, Write};

fn main() {
    let out = &mut BufWriter::new(stdout());
    let mut s = String::new();
    stdin().read_line(&mut s).expect("Failed to read input");
    let s = s.trim();

    let mut q = String::new();
    stdin()
        .read_line(&mut q)
        .expect("Failed to read number of queries");

    let q: u32 = q.trim().parse().unwrap();

    if q == 0 {
        writeln!(out, "No").ok();
    }

    let sorted_suffixes = suffix_array(&s);
    let mut saved_idx = HashMap::with_capacity(q as usize);
    let mut saved_results = HashMap::with_capacity(q as usize);

    for idx in 0..q {
        let mut i = String::new();

        stdin()
            .read_line(&mut i)
            .unwrap_or_else(|_| panic!("Failed to read input {}", idx));
        let i = i.trim();
        let result = saved_results.entry(i.to_string()).or_insert_with(|| {
            let sorted_suffixes = saved_idx
                .entry(i.len())
                .or_insert_with(|| transform(&sorted_suffixes, s, i.len()));

            match sorted_suffixes.binary_search(&i) {
                Ok(_) => "Yes",
                _ => "No",
            }
        });

        writeln!(out, "{}", result).ok();
    }
}

fn transform<'a>(prep: &[usize], s: &'a str, max_len: usize) -> Vec<&'a str> {
    let str_len = s.len();
    let mut result = Vec::with_capacity(prep.len());
    for &p in prep {
        if p + max_len > str_len {
            result.push(&s[p..])
        } else {
            result.push(&s[p..p + max_len])
        }
    }
    result
}

fn suffix_array(s: &str) -> Vec<usize> {
    let s = format!("{}{}", s, "$");
    let str_len = s.len();

    let (mut prep, mut class) = (vec![0; str_len], vec![0; str_len]);
    let mut arr: Vec<(char, usize)> = Vec::with_capacity(str_len);

    for (idx, ch) in s.chars().enumerate() {
        arr.push((ch, idx));
    }

    arr.sort_by(|a, b| a.0.cmp(&b.0));

    for (idx, a) in arr.iter().enumerate() {
        prep[idx] = a.1;
    }

    for (idx, _) in s.chars().enumerate().skip(1) {
        if arr[idx].0 == arr[idx - 1].0 {
            class[prep[idx]] = class[prep[idx - 1]];
        } else {
            class[prep[idx]] = class[prep[idx - 1]] + 1;
        }
    }

    let mut k = 0;
    while (1 << k) < str_len {
        prep = prep
            .iter()
            .map(|p| (p + str_len - (1 << k)) % str_len)
            .collect();

        count_sort(&mut prep, &class);

        let mut c_new = vec![0; str_len];
        for idx in 1..str_len {
            let prev = (
                class[prep[idx - 1]],
                class[(prep[idx - 1] + (1 << k)) % str_len],
            );
            let cur = (class[prep[idx]], class[(prep[idx] + (1 << k)) % str_len]);
            if cur == prev {
                c_new[prep[idx]] = c_new[prep[idx - 1]];
            } else {
                c_new[prep[idx]] = c_new[prep[idx - 1]] + 1;
            }
        }
        class = c_new;
        k += 1;
    }

    prep
}

fn count_sort(p: &mut Vec<usize>, c: &[i32]) {
    let n = &p.len();
    let mut cnt = vec![0; *n];
    for idx in 0..c.len() {
        cnt[c[idx] as usize] += 1;
    }
    let mut p_new = vec![0; *n];
    let mut pos = vec![0; *n];
    for idx in 1..*n {
        pos[idx] = pos[idx - 1] + cnt[idx - 1];
    }
    p.iter().for_each(|&item| {
        let i = c[item] as usize;
        p_new[pos[i]] = item;
        pos[i] += 1;
    });

    *p = p_new;
}


Get this bounty!!!

Leave a Reply

This site uses Akismet to reduce spam. Learn how your comment data is processed.