1 / 3
Oct 2023

Hello, I’m solving the MEANARR problem (https://www.spoj.com/problems/MEANARR/8).

I keep getting WA (e.g. on subission 32049468) and cannot find a test case where my program gives incorrect results.

May I ask for help? An example test case for which my code returns wrong answer would be helpful.

Source code (written in Rust):

use std::collections::HashSet;

fn read_line_to_u64_vec() -> Vec<i64> {
    let mut string = String::new();
    let _ = std::io::stdin().read_line(&mut string);
    return string.trim().split(' ').map(|s| s.parse().unwrap()).collect();
}

fn get_min_power_geq(number: usize) -> usize {

    let mut power_of_two = 1;
    while power_of_two < number {
        power_of_two = power_of_two * 2;
    }
    return power_of_two;
}

#[derive(Clone)]
#[derive(Debug)]
struct Node {
    pub left_range: i64,
    pub mid_range: i64,
    pub right_range: i64,

    pub count: u64
}

fn create_segment_tree(prefix_sums: &Vec<i64>) -> Vec<Node> {

    let mut prefix_sums_set: HashSet<i64> = HashSet::new();
    for number in prefix_sums.iter() {
        prefix_sums_set.insert(*number);
    }

    let segment_tree_len = 2 * get_min_power_geq(prefix_sums_set.len()) + 1;
    let max_int32 = 2147483647;
    let mut segment_tree = vec![Node{left_range: max_int32, right_range: max_int32, mid_range:max_int32, count: 0} ; segment_tree_len];

    let mut idx = segment_tree_len/2;
    let mut sorted_keys = prefix_sums_set.into_iter().collect::<Vec<i64>>();
    sorted_keys.sort();
    for key in sorted_keys {
        segment_tree[idx].left_range = key;
        segment_tree[idx].right_range = key;
        segment_tree[idx].mid_range = key;
        
        idx += 1;
    }

    
    for idx in (1..(segment_tree_len/2)).rev() {
        segment_tree[idx].left_range = segment_tree[2 * idx].left_range;
        segment_tree[idx].mid_range = segment_tree[2 * idx].right_range;
        segment_tree[idx].right_range = segment_tree[2 * idx + 1].right_range;
    }

    return segment_tree;
}

fn count_subsums_leq(segment_tree: &Vec<Node>, number: i64) -> u64 {
    let mut idx = 1;
    let mut count = 0;
     while segment_tree[idx].left_range != segment_tree[idx].right_range {

         if number > segment_tree[idx].mid_range {
             count += segment_tree[2 * idx].count;
             idx = 2 * idx + 1;
         }
         else {
             idx = 2 * idx;
         }
     }
     count += segment_tree[idx].count;
     return count;
}

fn insert_into_segment_tree(segment_tree: &mut Vec<Node>, number: i64) {
    let mut idx = 1;
    while segment_tree[idx].left_range != segment_tree[idx].right_range {
        segment_tree[idx].count += 1;
        if number > segment_tree[idx].mid_range {
            
            idx = 2 * idx + 1;
        }
        else {
            idx = 2 * idx;
        }
    }
    segment_tree[idx].count += 1;
}

fn main() {
    let mean = read_line_to_u64_vec()[1];
    let numbers = read_line_to_u64_vec();
    let prefix_sums = numbers.iter().scan(0, |sum, i| {*sum += i -mean; Some(*sum)}).collect::<Vec<_>>();

    let mut segment_tree = create_segment_tree(&prefix_sums);
    
    let mut count: u64 = 0;
    for number in prefix_sums.iter() {

        count += count_subsums_leq(&segment_tree, *number);
        if *number >= 0 {
            count += 1;
        }

        insert_into_segment_tree(&mut segment_tree, *number);
    }
    println!("{count}", count=count);
}
  • created

    Oct '23
  • last reply

    Oct '23
  • 2

    replies

  • 382

    views

  • 2

    users

  • 1

    link

I have a different answer for the following input:

10 1806598667
3458209057 848927043 4168239702 3656103956 3097075367 3427092634 431070554 19752024 3268890211 2528802850

Output:

45

Thanks for the test case. There was a tricky bug in an algorithm that showed up when numbers were big. I solved this problem.