#991093 (Rust) No.738 平らな農地

• 問題: No.738 平らな農地
• ユーザー: naut3
• 提出日時: 2024-06-28 13:05:49
• 言語: Rust (1.77.0 + proconio)
• 結果: AC
• 実行時間: 570 ms / 2,000 ms
• コード長: 12,123 bytes
• コンパイル時間: 14,060 ms
• コンパイル使用メモリ: 389,508 KB
• 実行使用メモリ: 31,776 KB
• 最終ジャッジ日時: 2024-06-28 13:06:23
• 合計ジャッジ時間: 31,989 ms
• ジャッジサーバーID （参考情報）: judge3 / judge1

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 5,248 KB
testcase_01	AC	1 ms; 5,376 KB
testcase_02	AC	1 ms; 5,376 KB
testcase_03	AC	1 ms; 5,376 KB
testcase_04	AC	1 ms; 5,376 KB
testcase_05	AC	11 ms; 5,376 KB
testcase_06	AC	13 ms; 5,376 KB
testcase_07	AC	9 ms; 5,376 KB
testcase_08	AC	4 ms; 5,376 KB
testcase_09	AC	3 ms; 5,376 KB
testcase_10	AC	2 ms; 5,376 KB
testcase_11	AC	6 ms; 5,376 KB
testcase_12	AC	2 ms; 5,376 KB
testcase_13	AC	11 ms; 5,376 KB
testcase_14	AC	4 ms; 5,376 KB
testcase_15	AC	167 ms; 28,920 KB
testcase_16	AC	298 ms; 28,368 KB
testcase_17	AC	347 ms; 28,364 KB
testcase_18	AC	192 ms; 27,556 KB
testcase_19	AC	433 ms; 29,740 KB
testcase_20	AC	196 ms; 29,184 KB
testcase_21	AC	227 ms; 30,192 KB
testcase_22	AC	175 ms; 28,684 KB
testcase_23	AC	383 ms; 29,276 KB
testcase_24	AC	224 ms; 28,448 KB
testcase_25	AC	3 ms; 5,376 KB
testcase_26	AC	3 ms; 5,376 KB
testcase_27	AC	2 ms; 5,376 KB
testcase_28	AC	2 ms; 5,376 KB
testcase_29	AC	4 ms; 5,376 KB
testcase_30	AC	4 ms; 5,376 KB
testcase_31	AC	2 ms; 5,376 KB
testcase_32	AC	3 ms; 5,376 KB
testcase_33	AC	3 ms; 5,376 KB
testcase_34	AC	3 ms; 5,376 KB
testcase_35	AC	2 ms; 5,376 KB
testcase_36	AC	2 ms; 5,376 KB
testcase_37	AC	3 ms; 5,376 KB
testcase_38	AC	3 ms; 5,376 KB
testcase_39	AC	3 ms; 5,376 KB
testcase_40	AC	3 ms; 5,376 KB
testcase_41	AC	4 ms; 5,376 KB
testcase_42	AC	3 ms; 5,376 KB
testcase_43	AC	3 ms; 5,376 KB
testcase_44	AC	3 ms; 5,376 KB
testcase_45	AC	154 ms; 30,940 KB
testcase_46	AC	260 ms; 28,068 KB
testcase_47	AC	152 ms; 29,488 KB
testcase_48	AC	210 ms; 28,172 KB
testcase_49	AC	210 ms; 29,008 KB
testcase_50	AC	205 ms; 29,156 KB
testcase_51	AC	167 ms; 30,264 KB
testcase_52	AC	149 ms; 28,236 KB
testcase_53	AC	153 ms; 28,900 KB
testcase_54	AC	161 ms; 30,556 KB
testcase_55	AC	196 ms; 30,536 KB
testcase_56	AC	266 ms; 30,140 KB
testcase_57	AC	243 ms; 27,992 KB
testcase_58	AC	150 ms; 29,776 KB
testcase_59	AC	237 ms; 30,900 KB
testcase_60	AC	252 ms; 30,488 KB
testcase_61	AC	235 ms; 29,932 KB
testcase_62	AC	141 ms; 27,876 KB
testcase_63	AC	173 ms; 30,124 KB
testcase_64	AC	159 ms; 29,748 KB
testcase_65	AC	288 ms; 27,900 KB
testcase_66	AC	300 ms; 29,704 KB
testcase_67	AC	196 ms; 28,916 KB
testcase_68	AC	290 ms; 30,332 KB
testcase_69	AC	513 ms; 29,864 KB
testcase_70	AC	365 ms; 30,996 KB
testcase_71	AC	193 ms; 29,060 KB
testcase_72	AC	330 ms; 28,824 KB
testcase_73	AC	292 ms; 29,476 KB
testcase_74	AC	373 ms; 29,620 KB
testcase_75	AC	445 ms; 28,224 KB
testcase_76	AC	327 ms; 30,856 KB
testcase_77	AC	511 ms; 28,732 KB
testcase_78	AC	341 ms; 31,440 KB
testcase_79	AC	327 ms; 31,776 KB
testcase_80	AC	233 ms; 29,912 KB
testcase_81	AC	284 ms; 28,792 KB
testcase_82	AC	431 ms; 30,156 KB
testcase_83	AC	200 ms; 30,088 KB
testcase_84	AC	194 ms; 29,716 KB
testcase_85	AC	570 ms; 31,768 KB
testcase_86	AC	295 ms; 28,392 KB
testcase_87	AC	57 ms; 30,060 KB
testcase_88	AC	55 ms; 28,736 KB
testcase_89	AC	1 ms; 5,376 KB
testcase_90	AC	1 ms; 5,376 KB
testcase_91	AC	1 ms; 5,376 KB

ソースコード

#![allow(non_snake_case, unused_imports)]
use proconio::{fastout, input, marker::*};
use wavelet_matrix::WaveletMatrix;

#[fastout]
fn main() {
    input! {
        N: usize, K: usize,
        A: [u64; N],
    }

    let wm = WaveletMatrix::from_weighted_own(&A, 30);

    let mut ans = u64::MAX;

    for i in 0..N {
        if i + K - 1 >= N {
            break;
        }

        let calc_ans = |j: usize| {
            let m = wm.quantile(i, i + K, j);
            let c1 = wm.range_freq(i, i + K, m);
            let c2 = (i + K - i) as u64 - c1;
            let s1 = wm.range_sum(i, i + K, m);
            let s2 = wm.sum(i, i + K) - s1;

            let mut ret = 0;
            ret += c1 * m - s1;
            ret += s2 - c2 * m;
            ret
        };

        if K % 2 == 1 {
            ans = std::cmp::min(ans, calc_ans(K / 2));
        } else {
            ans = std::cmp::min(ans, calc_ans(K / 2));
            ans = std::cmp::min(ans, calc_ans((K - 1) / 2));
        }
    }

    println!("{}", ans);
}

pub mod wavelet_matrix {
    pub struct WaveletMatrix<T> {
        bvs: Vec<BitVector>,
        length: usize,
        height: usize,
        cums: Vec<Vec<T>>,
    }

    impl WaveletMatrix<()> {
        /// 総和系クエリを利用しない場合のWavelet Matrixを構築する
        pub fn from(array: &[u64], height: usize) -> Self {
            let mut bvs = vec![];
            let mut array = array.to_vec();

            for i in (0..height).rev() {
                // i bit 目で安定ソートする
                let mut a0 = vec![];
                let mut a1 = vec![];
                let mut bv = vec![];

                for (j, &a) in array.iter().enumerate() {
                    if j % 64 == 0 {
                        bv.push(0);
                    }

                    bv[j / 64] |= ((a >> i) & 1) << (j % 64);

                    if (a >> i) & 1 == 1 {
                        a1.push(a);
                    } else {
                        a0.push(a);
                    }
                }

                bvs.push(BitVector::from(&bv));
                a0.append(&mut a1);

                array = a0;
            }

            bvs.reverse();

            Self {
                bvs,
                length: array.len(),
                height,
                cums: vec![],
            }
        }
    }

    impl<T> WaveletMatrix<T> {
        /// i 番目の要素の値を取得する
        pub fn access(&self, mut i: usize) -> u64 {
            // 上のbitから順番に位置を変更しながら走査すればよい
            let mut ret = 0u64;

            for j in (0..self.height).rev() {
                if self.bvs[j].access(i) {
                    ret |= 1 << j;

                    // 0の数 + 1の数
                    i = self.bvs[j].rank(i, true) as usize
                        + self.bvs[j].rank(self.length - 1, false) as usize
                        - 1;
                } else {
                    i = self.bvs[j].rank(i, false) as usize - 1;
                }
            }

            return ret;
        }

        /// [l, r) の中で k 番目に小さい値を求める (0 <= k)
        pub fn quantile(&self, mut l: usize, mut r: usize, mut k: usize) -> u64 {
            let mut ret = 0u64;

            for j in (0..self.height).rev() {
                let l0 = if l > 0 {
                    self.bvs[j].rank(l - 1, false)
                } else {
                    0
                };
                let r0 = if r > 0 {
                    self.bvs[j].rank(r - 1, false)
                } else {
                    0
                };

                if k as u32 + l0 < r0 {
                    l = l0 as usize;
                    r = r0 as usize;
                } else {
                    ret |= 1 << j;
                    k -= r0 as usize - l0 as usize;
                    let count_zeros = self.bvs[j].rank(self.length - 1, false);
                    l += (count_zeros - l0) as usize;
                    r += (count_zeros - r0) as usize;
                }
            }

            return ret;
        }

        /// [l, r) で upper 未満の要素の数を求める
        pub fn range_freq(&self, mut l: usize, mut r: usize, upper: u64) -> u64 {
            let mut ret = 0u64;

            for j in (0..self.height).rev() {
                let l0 = if l > 0 {
                    self.bvs[j].rank(l - 1, false)
                } else {
                    0
                };
                let r0 = if r > 0 {
                    self.bvs[j].rank(r - 1, false)
                } else {
                    0
                };

                if (upper >> j) & 1 == 1 {
                    ret += (r0 - l0) as u64;
                    let count_zeros = self.bvs[j].rank(self.length - 1, false);

                    l += (count_zeros - l0) as usize;
                    r += (count_zeros - r0) as usize;
                } else {
                    l = l0 as usize;
                    r = r0 as usize;
                }
            }

            ret
        }
    }

    impl WaveletMatrix<u64> {
        /// 自身の値を使った総和系クエリを利用する場合のWavelet Matrixを構築する
        pub fn from_weighted_own(array: &[u64], height: usize) -> Self {
            let mut bvs = vec![];
            let mut cums = vec![];
            let mut array = array.to_vec();

            for i in (0..height).rev() {
                // i bit 目で安定ソートする
                let mut a0 = vec![];
                let mut a1 = vec![];
                let mut bv = vec![];

                for (j, &a) in array.iter().enumerate() {
                    if j % 64 == 0 {
                        bv.push(0);
                    }

                    bv[j / 64] |= ((a >> i) & 1) << (j % 64);

                    if (a >> i) & 1 == 1 {
                        a1.push(a);
                    } else {
                        a0.push(a);
                    }
                }

                bvs.push(BitVector::from(&bv));
                a0.append(&mut a1);

                let mut cs = vec![0];

                for j in 0..array.len() {
                    cs.push(cs[j] + a0[j]);
                }

                cums.push(cs);
                array = a0;
            }

            bvs.reverse();
            cums.reverse();

            Self {
                bvs,
                length: array.len(),
                height,
                cums,
            }
        }
    }

    impl<T: Default + std::ops::Add<Output = T> + Clone + Copy> WaveletMatrix<T> {
        /// 自身の値を使わない総和系クエリを利用する場合のWavelet Matrixを構築する
        pub fn from_weighted(array: &[(u64, T)], height: usize) -> Self {
            let mut bvs = vec![];
            let mut cums = vec![];
            let mut array = array.to_vec();

            for i in (0..height).rev() {
                // i bit 目で安定ソートする
                let mut a0 = vec![];
                let mut a1 = vec![];
                let mut bv = vec![];

                for (j, &a) in array.iter().enumerate() {
                    if j % 64 == 0 {
                        bv.push(0);
                    }

                    bv[j / 64] |= ((a.0 >> i) & 1) << (j % 64);

                    if (a.0 >> i) & 1 == 1 {
                        a1.push(a);
                    } else {
                        a0.push(a);
                    }
                }

                bvs.push(BitVector::from(&bv));
                a0.append(&mut a1);

                let mut cs = vec![T::default()];

                for j in 0..array.len() {
                    cs.push(cs[j] + a0[j].1);
                }

                cums.push(cs);
                array = a0;
            }

            bvs.reverse();
            cums.reverse();

            Self {
                bvs,
                length: array.len(),
                height,
                cums,
            }
        }
    }

    impl<T: Default + std::ops::Add<Output = T> + Clone + Copy + std::ops::Sub<Output = T>>
        WaveletMatrix<T>
    {
        /// [l, r) で upper 未満の要素の総和を求める
        pub fn range_sum(&self, mut l: usize, mut r: usize, upper: u64) -> T {
            let mut ret = T::default();

            for j in (0..self.height).rev() {
                let l0 = if l > 0 {
                    self.bvs[j].rank(l - 1, false)
                } else {
                    0
                };
                let r0 = if r > 0 {
                    self.bvs[j].rank(r - 1, false)
                } else {
                    0
                };

                if (upper >> j) & 1 == 1 {
                    ret = ret + self.cums[j][r0 as usize] - self.cums[j][l0 as usize];
                    let count_zeros = self.bvs[j].rank(self.length - 1, false);
                    l += (count_zeros - l0) as usize;
                    r += (count_zeros - r0) as usize;
                } else {
                    l = l0 as usize;
                    r = r0 as usize;
                }
            }

            ret
        }

        /// [l, r) の要素の総和を求める
        pub fn sum(&self, mut l: usize, mut r: usize) -> T {
            let mut ret = T::default();

            for j in (0..self.height).rev() {
                let l0 = if l > 0 {
                    self.bvs[j].rank(l - 1, false)
                } else {
                    0
                };
                let r0 = if r > 0 {
                    self.bvs[j].rank(r - 1, false)
                } else {
                    0
                };

                ret = ret + self.cums[j][r0 as usize] - self.cums[j][l0 as usize];
                let count_zeros = self.bvs[j].rank(self.length - 1, false);
                l += (count_zeros - l0) as usize;
                r += (count_zeros - r0) as usize;
            }

            ret
        }
    }

    impl<T> std::fmt::Display for WaveletMatrix<T> {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            for i in 0..64 {
                let _ = writeln!(f, "{}", self.bvs[i]);
            }
            Ok(())
        }
    }

    pub struct BitVector {
        row: Vec<u64>,
        cs: Vec<u32>,
    }

    impl BitVector {
        pub fn from(array: &[u64]) -> Self {
            let mut cs = vec![0];

            for (i, a) in array.iter().enumerate() {
                cs.push(cs[i] + a.count_ones());
            }

            return Self {
                row: array.to_vec(),
                cs,
            };
        }

        /// i bit 目の値を返す (i は 0 以上)
        pub fn access(&self, i: usize) -> bool {
            (self.row[i / 64] >> (i % 64)) & 1 == 1
        }

        /// i bit 目までの b の数を数える (0 bit 目の存在に注意)
        pub fn rank(&self, i: usize, b: bool) -> u32 {
            if i % 64 == 63 {
                let c = self.cs[i / 64 + 1];

                if b {
                    c
                } else {
                    i as u32 + 1 - c
                }
            } else {
                let c = self.cs[i / 64];
                let dc = (self.row[i / 64] & ((1 << ((i % 64) + 1)) - 1)).count_ones();

                if b {
                    c + dc
                } else {
                    i as u32 + 1 - (c + dc)
                }
            }
        }

        /// とりあえず使わないので後回しにする
        #[allow(unused)]
        fn select(&self, i: usize, b: bool) -> usize {
            todo!()
        }
    }

    impl std::fmt::Display for BitVector {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            write!(
                f,
                "{}",
                self.row
                    .iter()
                    .map(|x| format!("{:>064b}", x).chars().rev().collect::<String>())
                    .collect::<Vec<_>>()
                    .join("")
            )
        }
    }
}