結果

問題 No.924 紲星
ユーザー sakikuroesakikuroe
提出日時 2024-05-02 11:13:00
言語 Rust
(1.77.0 + proconio)
結果
AC  
実行時間 1,578 ms / 4,000 ms
コード長 9,263 bytes
コンパイル時間 13,922 ms
コンパイル使用メモリ 378,292 KB
実行使用メモリ 129,868 KB
最終ジャッジ日時 2024-11-23 00:33:07
合計ジャッジ時間 28,628 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge2
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 1 ms
5,248 KB
testcase_01 AC 1 ms
5,248 KB
testcase_02 AC 2 ms
5,248 KB
testcase_03 AC 4 ms
5,248 KB
testcase_04 AC 3 ms
5,248 KB
testcase_05 AC 4 ms
5,248 KB
testcase_06 AC 4 ms
5,248 KB
testcase_07 AC 3 ms
5,248 KB
testcase_08 AC 1,578 ms
129,868 KB
testcase_09 AC 1,374 ms
129,868 KB
testcase_10 AC 1,407 ms
129,740 KB
testcase_11 AC 1,360 ms
129,740 KB
testcase_12 AC 1,559 ms
129,868 KB
testcase_13 AC 639 ms
52,296 KB
testcase_14 AC 610 ms
37,968 KB
testcase_15 AC 571 ms
41,636 KB
testcase_16 AC 475 ms
101,300 KB
testcase_17 AC 994 ms
59,024 KB
testcase_18 AC 1 ms
5,248 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

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

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

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

    pub fn len(&self) -> usize {
        self.accum_table.len() - 1
    }

    pub fn access(&self, i: usize) -> usize {
        self.accum_table[i + 1] - self.accum_table[i]
    }

    pub fn rank(&self, i: usize) -> usize {
        self.accum_table[i]
    }
}

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

impl WaveletMatrix {
    const WAVELET_MATRIX_HEIGHT: usize = 32;

    pub fn new(v: &[usize]) -> Self {
        let mut v = v.to_vec();
        let mut accum = vec![0];
        for i in 0..v.len() {
            accum.push(accum[i] + v[i]);
        }
        let mut bit_table = vec![];
        let mut accum_table = vec![];
        for i in (0..WaveletMatrix::WAVELET_MATRIX_HEIGHT).rev() {
            bit_table.push(WaveletMatrixRow::new(
                v.iter().map(|&x| (x >> i) & 1).collect(),
            ));
            accum_table.push(WaveletMatrixRow::new(
                v.iter()
                    .map(|&x| if (x >> i) & 1 == 0 { x } else { 0 })
                    .collect::<Vec<_>>(),
            ));

            // slow code
            // ```
            // v.sort_by_key(|&x| (x >> i) & 1);
            // ```
            v = v
                .iter()
                .filter(|&x| (x >> i) & 1 == 0)
                .chain(v.iter().filter(|&x| (x >> i) & 1 == 1))
                .cloned()
                .collect::<Vec<_>>();
        }

        WaveletMatrix {
            accum,
            bit_table,
            accum_table,
        }
    }

    /// Returns:
    ///     v[i]
    pub fn access(&self, mut i: usize) -> Option<usize> {
        if i >= self.bit_table[0].len() {
            return None;
        }

        let mut res = 0;

        for row in &self.bit_table {
            res <<= 1;
            res |= row.access(i);
            i = match row.access(i) {
                0 => i - row.rank(i),
                1 => row.len() - row.rank(row.len()) + row.rank(i),
                _ => {
                    unreachable!()
                }
            };
        }

        Some(res)
    }

    /// Returns:
    ///     v[l..r].into_iter().filter(|y| **y == x).count()
    pub fn rank(&self, mut l: usize, mut r: usize, x: usize) -> usize {
        for (i, row) in self.bit_table.iter().rev().enumerate().rev() {
            if (x >> i) & 1 == 0 {
                l = l - row.rank(l);
                r = r - row.rank(r);
            } else {
                l = row.rank(l) + (row.len() - row.rank(row.len()));
                r = row.rank(r) + (row.len() - row.rank(row.len()));
            }
        }

        r - l
    }

    /// 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 l >= r || r - l <= k {
            return None;
        }

        let mut res = 0;

        for (i, row) in self.bit_table.iter().rev().enumerate().rev() {
            let j = (r - l) - (row.rank(r) - row.rank(l));
            if j > k {
                l = l - row.rank(l);
                r = r - row.rank(r);
            } else {
                l = row.rank(l) + (row.len() - row.rank(row.len()));
                r = row.rank(r) + (row.len() - row.rank(row.len()));
                res |= 1 << i;
                k -= j;
            }
        }

        Some(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 + 1 + k) {
            None
        } else {
            self.get_kth_smallest(l, r, r - l - 1 - k)
        }
    }

    /// Returns:
    ///     {
    ///         use itertools::Itertools;
    ///         v[l..r].into_iter().sorted().take(k).sum::<usize>()
    ///     }
    pub fn get_sum_lower_k(&self, mut l: usize, mut r: usize, mut k: usize) -> usize {
        let mut res = 0;
        let mut kth = 0;

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

        res += k * kth;

        res
    }

    /// Returns:
    ///     {
    ///         use itertools::Itertools;
    ///         v[l..r].into_iter().sorted().rev().take(k).sum::<usize>()
    ///     }
    pub fn get_sum_upper_k(&self, l: usize, r: usize, k: usize) -> usize {
        self.accum[r] - self.accum[l] - self.get_sum_lower_k(l, r, r - l - k)
    }

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

    /// 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 {
        let mut res = 0;

        for (i, row) in self.bit_table.iter().rev().enumerate().rev() {
            if (upper >> i) & 1 == 0 {
                l = l - row.rank(l);
                r = r - row.rank(r);
            } else {
                res += (r - l) - (row.rank(r) - row.rank(l));
                l = row.rank(l) + (row.len() - row.rank(row.len()));
                r = row.rank(r) + (row.len() - row.rank(row.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 {
        r - l - self.count_less_than(l, r, lower)
    }

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

    /// Returns:
    ///     {
    ///         use itertools::Itertools;
    ///         v[l..r].into_iter().filter(|y| lower <= **y).sorted().nth(k)
    ///     }
    pub fn get_above_kth_smallest(
        &self,
        l: usize,
        r: usize,
        lower: usize,
        k: usize,
    ) -> Option<usize> {
        let cnt = self.count(l, r, 0, lower);
        if cnt + k >= r - l {
            None
        } else {
            Some(self.get_kth_smallest(l, r, cnt + k).unwrap())
        }
    }

    /// Returns:
    ///     v[l..r].into_iter().filter(|y| **y < upper).sorted().rev().nth(k)
    pub fn get_below_kth_largest(
        &self,
        l: usize,
        r: usize,
        upper: usize,
        k: usize,
    ) -> Option<usize> {
        let cnt = self.count(l, r, 0, upper);
        if cnt <= k {
            None
        } else {
            Some(self.get_kth_smallest(l, r, cnt - 1 - k).unwrap())
        }
    }

    /// Returns:
    ///     v[l..r].map(|y| y.abs_diff(x)).sum::<usize>()
    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 = self.count_more_than(l, r, x);
        x * c - self.get_sum_lower_k(l, r, c) + self.get_sum_upper_k(l, r, d) - x * d
    }
}

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

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

    let mut ans = vec![];

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

        if (r - l) % 2 == 1 {
            let mid = wm.get_kth_largest(l, r, (r - l) / 2).unwrap();
            ans.push(wm.get_sum_abs_diff(l, r, mid) / 2);
        } else {
            let mid0 = wm.get_kth_largest(l, r, (r - l) / 2 - 1).unwrap();
            let mid1 = wm.get_kth_largest(l, r, (r - l) / 2).unwrap();
            ans.push(wm.get_sum_abs_diff(l, r, (mid0 + mid1) / 2) / 2);
        }
    }

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