結果

問題 No.2977 Kth Xor Pair
コンテスト
ユーザー ikd
提出日時 2024-12-01 19:26:45
言語 Rust
(1.83.0 + proconio)
結果
MLE  
実行時間 -
コード長 3,729 bytes
コンパイル時間 15,352 ms
コンパイル使用メモリ 380,720 KB
実行使用メモリ 821,668 KB
最終ジャッジ日時 2024-12-01 19:28:25
合計ジャッジ時間 92,734 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1 MLE * 1
other AC * 5 TLE * 4 MLE * 25
権限があれば一括ダウンロードができます

ソースコード

diff # 

use proconio::input;

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

    let mut trie = vec![BinaryTrie::new(); n];
    for i in 0..n {
        for j in 0..i {
            trie[i].add(a[j]);
        }
    }

    // ● xor ▲ ≦ x を満たす組の個数
    let count_leq = |x: u32| -> usize {
        let mut res = 0;
        for i in 0..n {
            res += trie[i].query(a[i], x);
        }
        res
    };

    if count_leq(0) >= k {
        println!("0");
        return;
    }

    // ● xor ▲ ≦ ok が k 個以上ある
    let mut ng = 0;
    let mut ok = u32::MAX / 2;
    while ok - ng > 1 {
        let mid = (ok + ng) / 2;
        if count_leq(mid) >= k {
            ok = mid;
        } else {
            ng = mid;
        }
    }

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

#[derive(Debug, Clone)]
struct Node {
    size: usize,
    children: [Option<Box<Node>>; 2],
}

impl Node {
    fn new() -> Self {
        Self {
            size: 0,
            children: [None, None],
        }
    }
}

#[derive(Debug, Clone)]
struct BinaryTrie {
    root: Node,
}

impl BinaryTrie {
    fn new() -> Self {
        Self { root: Node::new() }
    }

    fn add(&mut self, value: u32) {
        let mut node = &mut self.root;
        node.size += 1;
        for i in (0..u32::BITS).rev() {
            let index = (value >> i & 1) as usize;
            node = node.children[index].get_or_insert_with(|| Box::new(Node::new()));
            node.size += 1;
        }
    }

    fn query(&self, a: u32, x: u32) -> usize {
        // fn f(node: &Node, k: u32, a: u32, x: u32) -> usize {
        //     if k == 0 {
        //         let a_k = (a >> k & 1) as usize;
        //         if x >> k & 1 == 0 {
        //             match &node.children[a_k] {
        //                 Some(n) => n.size,
        //                 None => 0,
        //             }
        //         } else {
        //             match &node.children {
        //                 [Some(n), Some(m)] => n.size + m.size,
        //                 [Some(n), None] => n.size,
        //                 [None, Some(n)] => n.size,
        //                 [None, None] => 0,
        //             }
        //         }
        //     } else {
        //         let a_k = (a >> k & 1) as usize;
        //         if x >> k & 1 == 0 {
        //             match &node.children[a_k] {
        //                 Some(n) => f(n, k - 1, a, x),
        //                 None => 0,
        //             }
        //         } else {
        //             let c = match &node.children[a_k] {
        //                 Some(n) => n.size,
        //                 None => 0,
        //             };
        //             let d = match &node.children[a_k ^ 1] {
        //                 Some(n) => f(n, k - 1, a, x),
        //                 None => 0,
        //             };
        //             c + d
        //         }
        //     }
        // }
        // f(&self.root, u32::BITS - 1, a, x)

        let mut node = &self.root;
        let mut res = 0;
        for k in (0..u32::BITS).rev() {
            let a_k = (a >> k & 1) as usize;
            if x >> k & 1 == 0 {
                match &node.children[a_k] {
                    Some(n) => node = n,
                    None => return res,
                }
            } else {
                res += match &node.children[a_k] {
                    Some(n) => n.size,
                    None => 0,
                };
                match &node.children[a_k ^ 1] {
                    Some(n) => node = n,
                    None => return res,
                }
            }
        }
        res + node.size
    }
}
0