#978606 (Rust) No.924 紲星

• 問題: No.924 紲星
• ユーザー: sakikuroe
• 提出日時: 2024-05-02 14:15:15
• 言語: Rust (1.77.0 + proconio)
• 結果: AC
• 実行時間: 537 ms / 4,000 ms
• コード長: 9,259 bytes
• コンパイル時間: 14,649 ms
• コンパイル使用メモリ: 379,948 KB
• 実行使用メモリ: 84,500 KB
• 最終ジャッジ日時: 2024-11-23 04:39:52
• 合計ジャッジ時間: 20,802 ms
• ジャッジサーバーID （参考情報）: judge4 / judge1

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 5,248 KB
testcase_01	AC	1 ms; 5,248 KB
testcase_02	AC	1 ms; 5,248 KB
testcase_03	AC	2 ms; 5,248 KB
testcase_04	AC	1 ms; 5,248 KB
testcase_05	AC	2 ms; 5,248 KB
testcase_06	AC	2 ms; 5,248 KB
testcase_07	AC	2 ms; 5,248 KB
testcase_08	AC	529 ms; 84,500 KB
testcase_09	AC	509 ms; 84,376 KB
testcase_10	AC	537 ms; 84,380 KB
testcase_11	AC	533 ms; 84,372 KB
testcase_12	AC	525 ms; 84,372 KB
testcase_13	AC	224 ms; 35,808 KB
testcase_14	AC	196 ms; 27,128 KB
testcase_15	AC	197 ms; 29,388 KB
testcase_16	AC	246 ms; 62,296 KB
testcase_17	AC	342 ms; 42,120 KB
testcase_18	AC	1 ms; 5,248 KB

ソースコード

use proconio::{input, marker::Usize1};

#[derive(Clone)]
pub struct AccumulateVector {
    accum_table: Vec<usize>,
}

impl AccumulateVector {
    pub fn new(v: &[usize]) -> Self {
        let mut accum_table = vec![0; v.len() + 1];
        for i in 0..v.len() {
            accum_table[i + 1] = accum_table[i] + v[i];
        }
        AccumulateVector { accum_table }
    }

    /// Returns:
    ///     lenght of v
    pub fn len(&self) -> usize {
        self.accum_table.len() - 1
    }

    /// Returns:
    ///     sum(v[0..i])
    pub fn rank(&self, i: usize) -> usize {
        self.accum_table[i]
    }
}

#[derive(Clone)]
pub struct WaveletMatrix {
    bit_table: Vec<AccumulateVector>,
    sorted_deduped_v: Vec<usize>,
    accum_table: Vec<AccumulateVector>,
    accum_v: AccumulateVector,
}

impl WaveletMatrix {
    const WAVELET_MATRIX_HEIGHT: usize = 18;

    pub fn new(v: &[usize]) -> Self {
        let mut sorted_deduped_v = v.to_vec();
        sorted_deduped_v.sort_unstable();
        sorted_deduped_v.dedup();
        let mut compress = v
            .iter()
            .map(|&x| sorted_deduped_v.partition_point(|&y| y < x))
            .collect::<Vec<_>>();
        let mut bit_table = vec![];
        let mut accum_table = vec![];
        for i in (0..WaveletMatrix::WAVELET_MATRIX_HEIGHT).rev() {
            bit_table.push(AccumulateVector::new(
                &compress.iter().map(|&x| (x >> i) & 1).collect::<Vec<_>>(),
            ));
            accum_table.push(AccumulateVector::new(
                &compress
                    .iter()
                    .map(|&x| {
                        if (x >> i) & 1 == 0 {
                            sorted_deduped_v[x]
                        } else {
                            0
                        }
                    })
                    .collect::<Vec<_>>(),
            ));
            compress = compress
                .iter()
                .filter(|&x| (x >> i) & 1 == 0)
                .chain(compress.iter().filter(|&x| (x >> i) & 1 == 1))
                .cloned()
                .collect::<Vec<_>>();
        }

        let accum_v = AccumulateVector::new(v);

        WaveletMatrix {
            bit_table,
            sorted_deduped_v,
            accum_table,
            accum_v,
        }
    }

    /// Returns:
    ///     v[l..r].into_iter().filter(|y| y < upper).count()
    pub fn count_less_than(&self, mut l: usize, mut r: usize, upper: usize) -> usize {
        if r <= l {
            return 0;
        }

        let upper = self.sorted_deduped_v.partition_point(|&x| x < upper);
        let mut res = 0;

        for (i, bit) in (0..Self::WAVELET_MATRIX_HEIGHT)
            .rev()
            .zip(self.bit_table.iter())
        {
            if (upper >> i) & 1 == 0 {
                l = l - bit.rank(l);
                r = r - bit.rank(r);
            } else {
                res += (r - l) - (bit.rank(r) - bit.rank(l));
                l = bit.rank(l) + (bit.len() - bit.rank(bit.len()));
                r = bit.rank(r) + (bit.len() - bit.rank(bit.len()));
            }
        }

        res
    }

    /// Returns:
    ///     v[l..r].into_iter().filter(|y| lower <= y ).count()
    pub fn count_more_than(&self, l: usize, r: usize, lower: usize) -> usize {
        if r <= l {
            return 0;
        }

        r - l - self.count_less_than(l, r, lower)
    }

    /// Returns:
    ///     {
    ///         use itertools::Itertools;
    ///         v[l..r].into_iter().sorted().nth(k)
    ///     }
    pub fn get_kth_smallest(&self, mut l: usize, mut r: usize, mut k: usize) -> Option<usize> {
        if r <= l || r - l <= k {
            return None;
        }

        let mut res = 0;

        for (i, bit) in (0..Self::WAVELET_MATRIX_HEIGHT)
            .rev()
            .zip(self.bit_table.iter())
        {
            let j = (r - l) - (bit.rank(r) - bit.rank(l));
            if k < j {
                l = l - bit.rank(l);
                r = r - bit.rank(r);
            } else {
                l = bit.rank(l) + (bit.len() - bit.rank(bit.len()));
                r = bit.rank(r) + (bit.len() - bit.rank(bit.len()));
                res |= 1 << i;
                k -= j;
            }
        }

        Some(self.sorted_deduped_v[res])
    }

    /// Returns:
    ///     {
    ///         use itertools::Itertools;
    ///         v[l..r].into_iter().sorted().rev().nth(k)
    ///     }
    pub fn get_kth_largest(&self, l: usize, r: usize, k: usize) -> Option<usize> {
        if r <= l || r - l <= k {
            return None;
        }

        self.get_kth_smallest(l, r, r - l - 1 - k)
    }

    /// Returns:
    ///     v[l..r].into_iter().filter(|y| lower <= y < upper).count()
    pub fn count(&self, l: usize, r: usize, lower: usize, upper: usize) -> usize {
        if r <= l {
            return 0;
        }

        self.count_less_than(l, r, upper) - self.count_less_than(l, r, lower)
    }

    /// Returns:
    ///     {
    ///         use itertools::Itertools;
    ///         v[l..r].into_iter().filter(|x| x < upper).sum()
    ///     }
    pub fn get_sum_less_than(&self, mut l: usize, mut r: usize, upper: usize) -> usize {
        if r <= l {
            return 0;
        }

        let upper = self.sorted_deduped_v.partition_point(|&x| x < upper);

        let mut res = 0;

        for (i, (bit, accum)) in (0..Self::WAVELET_MATRIX_HEIGHT)
            .rev()
            .zip(self.bit_table.iter().zip(self.accum_table.iter()))
        {
            if (upper >> i) & 1 == 0 {
                l = l - bit.rank(l);
                r = r - bit.rank(r);
            } else {
                res += accum.rank(r) - accum.rank(l);
                l = bit.rank(l) + (bit.len() - bit.rank(bit.len()));
                r = bit.rank(r) + (bit.len() - bit.rank(bit.len()));
            }
        }

        res
    }

    /// Returns:
    ///     {
    ///         use itertools::Itertools;
    ///         v[l..r].into_iter().filter(|x| lower <= x).sum()
    ///     }
    pub fn get_sum_more_than(&self, l: usize, r: usize, lower: usize) -> usize {
        if r <= l {
            return 0;
        }

        self.accum_v.rank(r) - self.accum_v.rank(l) - self.get_sum_less_than(l, r, lower)
    }

    /// Returns:
    ///     {
    ///         use itertools::Itertools;
    ///         v[l..r].into_iter().sorted().take(k).sum()
    ///     }
    pub fn get_sum_k_smallest(&self, mut l: usize, mut r: usize, mut k: usize) -> Option<usize> {
        if r <= l || r - l < k {
            return None;
        }

        let mut res = 0;
        let mut kth = 0;

        for (i, (bit, sum)) in (0..Self::WAVELET_MATRIX_HEIGHT)
            .rev()
            .zip(self.bit_table.iter().zip(self.accum_table.iter()))
        {
            let j = (r - l) - (bit.rank(r) - bit.rank(l));
            if k < j {
                l = l - bit.rank(l);
                r = r - bit.rank(r);
            } else {
                res += sum.rank(r);
                res -= sum.rank(l);
                l = bit.rank(l) + (bit.len() - bit.rank(bit.len()));
                r = bit.rank(r) + (bit.len() - bit.rank(bit.len()));
                kth |= 1 << i;
                k -= j;
            }
        }

        res += k * kth;

        Some(res)
    }

    /// Returns:
    ///     {
    ///         use itertools::Itertools;
    ///         v[l..r].into_iter().sorted().rev().take(k).sum()
    ///     }
    pub fn get_sum_k_largest(&self, l: usize, r: usize, k: usize) -> Option<usize> {
        if r <= l || r - l < k {
            return None;
        }

        Some(
            self.accum_v.rank(r)
                - self.accum_v.rank(l)
                - self.get_sum_k_smallest(l, r, r - l - k).unwrap(),
        )
    }

    /// Returns:
    ///     {
    ///         use itertools::Itertools;
    ///         v[l..r].into_iter().filter(|x| lower <= x < upper).sum()
    ///     }
    pub fn get_sum(&self, l: usize, r: usize, lower: usize, upper: usize) -> usize {
        if r <= l {
            return 0;
        }

        self.get_sum_less_than(l, r, upper) - self.get_sum_less_than(l, r, lower)
    }

    /// Returns:
    ///     v[l..r].map(|y| y.abs_diff(x)).sum()
    pub fn get_sum_abs_diff(&self, l: usize, r: usize, x: usize) -> usize {
        let c = self.count_less_than(l, r, x);
        let d = r - l - c;
        let sum_less_than = self.get_sum_less_than(l, r, x);
        let sum_more_than: usize = self.accum_v.rank(r) - self.accum_v.rank(l) - sum_less_than;
        (x * c - sum_less_than) + sum_more_than - x * d
    }
}

fn main() {
    input! {
        n: usize, q: usize,
        a: [isize; n],
    }

    let wm = WaveletMatrix::new(
        &a.iter()
            .cloned()
            .map(|x| (x + 1000000000) as usize)
            .collect::<Vec<_>>(),
    );

    let mut ans = vec![];

    for _ in 0..q {
        input! {
            l: Usize1, r: usize,
        }

        let mid = wm.get_kth_largest(l, r, (r - l) / 2).unwrap();
        ans.push(wm.get_sum_abs_diff(l, r, mid));
    }

    println!(
        "{}",
        ans.into_iter()
            .map(|x| x.to_string())
            .collect::<Vec<_>>()
            .join("\n")
    );
}