結果

問題 No.2761 Substitute and Search
コンテスト
ユーザー magurofly
提出日時 2024-04-29 22:41:52
言語 Rust
(1.83.0 + proconio)
結果
MLE  
実行時間 -
コード長 3,965 bytes
コンパイル時間 14,115 ms
コンパイル使用メモリ 379,620 KB
実行使用メモリ 1,631,128 KB
最終ジャッジ日時 2024-11-29 18:09:23
合計ジャッジ時間 33,099 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 6 MLE * 7
権限があれば一括ダウンロードができます
コンパイルメッセージ
warning: unused import: `std::collections::*`
 --> src/main.rs:2:5
  |
2 | use std::collections::*;
  |     ^^^^^^^^^^^^^^^^^^^
  |
  = note: `#[warn(unused_imports)]` on by default

warning: variable `N` should have a snake case name
 --> src/main.rs:6:9
  |
6 |         N: usize, L: usize, Q: usize,
  |         ^ help: convert the identifier to snake case: `n`
  |
  = note: `#[warn(non_snake_case)]` on by default

warning: variable `L` should have a snake case name
 --> src/main.rs:6:19
  |
6 |         N: usize, L: usize, Q: usize,
  |                   ^ help: convert the identifier to snake case: `l`

warning: variable `Q` should have a snake case name
 --> src/main.rs:6:29
  |
6 |         N: usize, L: usize, Q: usize,
  |                             ^ help: convert the identifier to snake case: `q`

warning: variable `S` should have a snake case name
 --> src/main.rs:7:9
  |
7 |         S: [Chars; N],
  |         ^ help: convert the identifier to snake case (notice the capitalization): `s`

ソースコード

diff # 

use proconio::{input, marker::{Usize1, Chars}};
use std::collections::*;

fn main() {
    input! {
        N: usize, L: usize, Q: usize,
        S: [Chars; N],
    }

    let mut trie = Trie::new(L);
    for s in &S {
        trie.add(&lower_alphabets_to_usizes(&s), 1);
    }

    for _ in 0 .. Q {
        input!(qtype: usize);
        match qtype {
            1 => {
                input!(k: Usize1, c: char, d: char);
                let c = la2i(c);
                let d = la2i(d);
                trie.substitute(k, c, d);
            }
            2 => {
                input!(t: Chars);
                let t = lower_alphabets_to_usizes(&t);
                println!("{}", trie.search(&t));
            }
            _ => unreachable!()
        }
    }
}

fn la2i(c: char) -> usize {
    (c as u8 - b'a') as usize
}

fn lower_alphabets_to_usizes(chars: &[char]) -> Vec<usize> {
    chars.iter().map(|&c| la2i(c) ).collect()
}

struct Trie {
    nodes: Vec<TrieNode>,
    by_index_key: Vec<Vec<Vec<usize>>>,
}
impl Trie {
    pub fn new(l: usize) -> Self {
        Self {
            nodes: vec![TrieNode::new()],
            by_index_key: vec![vec![vec![]; 26]; l],
        }
    }

    pub fn add(&mut self, s: &[usize], add: i64) {
        let mut current = 0;
        self.nodes[current].sum += add;
        for i in 0 .. s.len() {
            let c = s[i];
            if self.nodes[current].children[c].is_none() {
                let id = self.nodes.len();
                let mut node = TrieNode::new();
                node.parent = Some(current);
                self.nodes.push(node);
                self.nodes[current].children[c] = Some(id);
                self.by_index_key[i][c].push(id);
            }
            current = self.nodes[current].children[c].unwrap();
            self.nodes[current].sum += add;
        }
    }

    pub fn substitute(&mut self, k: usize, from: usize, to: usize) {
        let src = std::mem::take(&mut self.by_index_key[k][from]);
        let mut merge_stack = vec![];
        for &id in &src {
            if self.nodes[id].is_removed {
                continue;
            }
            let Some(parent) = self.nodes[id].parent else { continue };
            if let Some(merged) = self.nodes[parent].children[to] {
                self.nodes[id].is_removed = true;
                merge_stack.push((id, merged, k));
            } else {
                self.nodes[parent].children[to] = Some(id);
                self.by_index_key[k][to].push(id);
                self.nodes[id].parent = Some(parent);
            }
            self.nodes[parent].children[from] = None;
        }
        while let Some((src, dst, k)) = merge_stack.pop() {
            self.nodes[dst].sum += self.nodes[src].sum;
            for c in 0 .. 26 {
                if let Some(src_child) = self.nodes[src].children[c] {
                    if let Some(dst_child) = self.nodes[dst].children[c] {
                        self.nodes[src_child].is_removed = true;
                        merge_stack.push((src_child, dst_child, k + 1));
                    } else {
                        self.nodes[dst].children[c] = Some(src_child);
                        self.nodes[src_child].parent = Some(dst);
                    }
                }
            }
        }
    }

    pub fn search(&self, s: &[usize]) -> i64 {
        let mut cur = 0;
        for i in 0 .. s.len() {
            let c = s[i];
            if let Some(child) = self.nodes[cur].children[c] {
                cur = child;
            } else {
                return 0;
            }
        }
        self.nodes[cur].sum
    }
}

struct TrieNode {
    pub is_removed: bool,
    pub parent: Option<usize>,
    pub sum: i64,
    pub children: Vec<Option<usize>>,
}
impl TrieNode {
    pub fn new() -> Self {
        Self {
            is_removed: false,
            parent: None,
            sum: 0,
            children: vec![None; 26],
        }
    }
}
0