算法编程题系列-原子的数量

原题描述

LeetCode726 原子的数量:给定一个形如"K4(ON(SO3)2)2"的化学表达式，返回每种原子的数量，比如这个例子里面返回"K4N2O14S4"，标识原子K,N,O,S分别有4， 2， 14， 4个，按照原子名称的字典序返回。

思路简述

一道不算难得难题，笔者认为难在如何优雅简洁地实现。比较难处理的就是括号，可以采用两种方式处理，第一种是递归，即括号里面的又是一个化学表达式，第二种是用栈，可以认为是在模拟递归，遇到括号，就往栈中压入一个空的计数哈希表，进行后续的计数。

代码实现

package codes

import (
	"fmt"
	"sort"
)

// 化学式定义
type Formula struct {
	atomCounts map[string]int
}

// NewFormula 通过传入的化学式构造Formula
func NewFormula(formula string) *Formula {
	f := &Formula{atomCounts: make(map[string]int)}
	n := len(formula)
	i := 0
	for i < n {
		if formula[i] >= 'A' && formula[i] <= 'Z' {
			atom := ParseString(formula, &i)
			if i < n && formula[i] >= '1' && formula[i] <= '9' {
				count := ParseNum(formula, &i)
				f.AddAtom(atom, count)
			} else {
				f.AddAtom(atom, 1)
			}
		} else if formula[i] == '(' {
			exp := ParseExp(formula, &i)
			f1 := NewFormula(exp)
			if i < n && formula[i] >= '1' && formula[i] <= '9' {
				count := ParseNum(formula, &i)
				f1.Multi(count)
			} 
			f.MergeFormula(f1)
		}
	}
	return f
}

// AddAtom 给atom原子的计数增加count
func (f *Formula) AddAtom(atom string, count int) {
	f.atomCounts[atom] += count
}

// Multi 给每个原子的计数乘以multi
func (f *Formula) Multi(multi int) {
	keys := make([]string, len(f.atomCounts))
	for atom := range f.atomCounts {
		keys = append(keys, atom)
	}
	for index := range keys {
		f.atomCounts[keys[index]] *= multi
	}
}

// MergeFormula 合并两个Formula
func (f *Formula) MergeFormula(other *Formula) {
	for atom, count := range other.atomCounts {
		f.atomCounts[atom] += count
	}
}

// String 返回Formula的字符串表达形式
func (f *Formula) String() string {
	keys := make([]string, 0)
	for atom := range f.atomCounts {
		keys = append(keys, atom)
	}
	sort.Strings(keys)   // 按照字典序排序
	var ret string
	for index := range keys {
		ret += keys[index]
		count := f.atomCounts[keys[index]]
		if count <= 1 {
			continue
		}
		ret += fmt.Sprintf("%d", count)
	}
	return ret
}


func countOfAtoms(formula string) string {
    f := NewFormula(formula)
	return f.String()
}


// ParseString 从exp的pos位置提取一个字符串，除了第一个为大写字符，区域均为小写字母
func ParseString(exp string, pos *int) string {
	j := *pos + 1
	for j < len(exp) && exp[j] >= 'a' && exp[j] <= 'z' {
		j++
	}
	ret := exp[*pos: j]
	*pos = j
	return ret
}


// ParseNum 从exp的pos位置开始提取一个数字
func ParseNum(exp string, pos *int) int {
	num := 0
	for *pos < len(exp) && exp[*pos] >= '0' && exp[*pos] <= '9' {
		num = num * 10 + int(exp[*pos] - '0')
		(*pos)++
	}
	return num
}

// ParseExp 提取括号内的一个表达式，用于递归处理
func  ParseExp(exp string, pos *int) string {
	i := *pos
	j := i + 1
	lbCount := 1   // 记录左括号的相对数量
	for j < len(exp) && (exp[j] != ')' || lbCount != 1) {
		if exp[j] == '(' {
			lbCount++
		} else if exp[j] == ')' {
			lbCount--
		}
		j++
	}
	*pos = j + 1
	return exp[i + 1: j]
}

复杂度分析

• 时间复杂度：$O(m\ast n)$，其中m为原子的类别数，n为化学表示式字符串的长度
• 空间复杂度：$O(m)$