Rust | Abstract Syntax Tree
Approach\n\nWrite a parser, construct an AST from the given string with that parser, then traverse the AST to count the number of occurences of each element.\n\
# Approach\n\nWrite a parser, construct an AST from the given string with that parser, then traverse the AST to count the number of occurences of each element.\n\n# Code\n```\nconst BASE: u32 = 10;\n\n#[derive(Debug, Clone)]\npub struct AtomNode(Vec<char>);\n\n#[derive(Debug, Clone)]\npub struct CountNode(usize);\n\n#[derive(Debug, Clone)]\npub enum ContainerNode {\n    Atom(AtomNode),\n    Formula(FormulaNode),\n}\n\n#[derive(Debug, Clone)]\npub struct FormulaItemNode(\n    ContainerNode,\n    Option<CountNode>,\n);\n\n#[derive(Debug, Clone)]\npub struct FormulaNode(\n    Vec<FormulaItemNode>\n);\n\npub trait ASTNode: Sized {\n    type Error;\n\n    fn parse(tokens: Vec<char>) -> Result<(Self, Vec<char>), Self::Error>;\n}\n\nimpl ASTNode for AtomNode {\n    type Error = &\'static str;\n\n    fn parse(mut tokens: Vec<char>) -> Result<(Self, Vec<char>), Self::Error> {\n        let mut ret = vec![];\n        \n        while let Some(token) = tokens.pop() {\n            if !token.is_alphabetic() {\n                return Err("AtomNode: illegal character");\n            }\n\n            ret.push(token);\n            if token.is_uppercase() {\n                break; // end of atomic symbol\n            }\n        }\n        ret.reverse();\n\n        Ok((Self(ret), tokens))\n    }\n}\n\nimpl ASTNode for CountNode {\n    type Error = &\'static str;\n\n    fn parse(mut tokens: Vec<char>) -> Result<(Self, Vec<char>), Self::Error> {\n        if !tokens.last()\n            .map(|&c| c.is_digit(BASE))\n            .unwrap_or(false) {\n            return Err("CountNode: no digit to collect");\n        }\n        \n        let mut ret = u32::MIN;\n\n        let mut base_cur = 1u32;\n        while let Some(token) = tokens.pop() {\n            if let Some(digit) = token.to_digit(BASE) {\n                ret += digit * base_cur;\n                base_cur *= BASE;\n            } else {\n                tokens.push(token);\n                break;\n            }\n        }\n\n        Ok((Self(ret as usize), tokens))\n    }\n}\n\nimpl ASTNode for ContainerNode {\n    type Error = &\'static str;\n\n    fn parse(mut tokens: Vec<char>) -> Result<(Self, Vec<char>), Self::Error> {\n        let c = if let Some(&inner) = tokens.last() {\n            inner\n        } else {\n            return Err("ContainerNode: no char to collect");\n        };\n\n        if c != \')\' {\n            let (inner, tokens) = AtomNode::parse(tokens)?;\n            return Ok((Self::Atom(inner), tokens));\n        }\n\n        tokens.pop();\n        let (inner, mut tokens) = FormulaNode::parse(tokens)?;\n        if tokens.pop().unwrap_or(\'X\') != \'(\' {\n            return Err("ContainerNode: mismatched brackets");\n        }\n\n        Ok((Self::Formula(inner), tokens))\n    }\n}\n\nimpl ASTNode for FormulaItemNode {\n    type Error = &\'static str;\n\n    fn parse(mut tokens: Vec<char>) -> Result<(Self, Vec<char>), Self::Error> {\n        let token = if let Some(&inner) = tokens.last() {\n            inner\n        } else {\n            return Err("FormulaItemNode: no char to collect");\n        };\n        \n        let mut count_node = None;\n        if token.is_digit(BASE) {\n            let (count_node_next, tokens_next) = CountNode::parse(tokens)?;\n            tokens = tokens_next;\n            count_node = Some(count_node_next);\n        }\n\n        let token = if let Some(&inner) = tokens.last() {\n            inner\n        } else {\n            return Err("FormulaItemNode: e0 must exist");\n        };\n\n        let (container_node, tokens) = ContainerNode::parse(tokens)?;\n        Ok((Self(container_node, count_node), tokens))\n    }\n}\n\nimpl ASTNode for FormulaNode {\n    type Error = &\'static str;\n\n    fn parse(mut tokens: Vec<char>) -> Result<(Self, Vec<char>), Self::Error> {\n        let mut ret = vec![];\n        \n        while let Some(&token) = tokens.last() {\n            if token == \'(\' {\n                break; // TODO: exclusion\n            }\n\n            let (item_node, tokens_next) = FormulaItemNode::parse(tokens)?;\n            tokens = tokens_next;\n\n            ret.push(item_node);\n        }\n        ret.reverse();\n\n        Ok((Self(ret), tokens))\n    }\n}\n\n#[derive(Debug, Clone)]\npub struct AST {\n    root: FormulaNode,\n}\n\nuse std::convert::TryFrom;\nimpl TryFrom<Vec<char>> for AST {\n    type Error = &\'static str;\n\n    fn try_from(tokens: Vec<char>) -> Result<Self, Self::Error> {\n        let (root, _) = FormulaNode::parse(tokens)?;\n        Ok(Self { root })\n    }\n}\n\nimpl FormulaNode {\n    pub fn stat(\n        self,\n        multiplier: usize,\n        tree: &mut BTreeMap<Vec<char>, usize>\n    ) {\n        for item in self.0 {\n            let FormulaItemNode(container_node, count_node) = item;\n            let count = count_node.map(|n| n.0).unwrap_or(1);\n            let count = count * multiplier;\n            match container_node {\n                ContainerNode::Formula(self_next) => {\n                    self_next.stat(count, tree);\n                },\n                ContainerNode::Atom(AtomNode(atom)) => {\n                    let count_prev = tree.remove(&atom).unwrap_or(usize::MIN);\n                    tree.insert(atom, count + count_prev);\n                },\n            }\n        }\n    }\n}\n\nuse std::collections::BTreeMap;\nimpl AST {\n    pub fn stat(self) -> Vec<(Vec<char>, usize)> {\n        let root = self.root;\n        \n        let mut tree = BTreeMap::new();\n        root.stat(1, &mut tree);\n        \n        tree.into_iter().collect::<Vec<_>>()\n    }\n}\n\nimpl Solution {\n    fn preproc(s: String) -> Vec<char> {\n        s.chars().collect::<Vec<_>>()\n    }\n\n    pub fn count_of_atoms(formula: String) -> String {\n        let tokens = Self::preproc(formula);\n\n        let ast = AST::try_from(tokens).unwrap();\n        // println!("ast: {:?}", ast);\n\n        let ret = ast.stat().into_iter()\n            .map(|(mut atom, count)| {\n                if count > 1 {\n                    let count_str = count.to_string().chars().collect::<Vec<_>>();\n                    atom.extend(count_str);\n                }\n                atom.into_iter().collect::<String>()\n            })\n            .collect::<Vec<_>>();\n\n        ret.join("")\n    }\n}\n```
['Rust']