Given a chemical formula (given as a string), return the count of each atom.

An atomic element always starts with an uppercase character, then zero or more lowercase letters, representing the name.

1 or more digits representing the count of that element may follow if the count is greater than 1.  If the count is 1, no digits will follow.  For example, H2O and H2O2 are possible, but H1O2 is impossible.

Two formulas concatenated together produce another formula.  For example, H2O2He3Mg4 is also a formula.  

A formula placed in parentheses, and a count (optionally added) is also a formula.  For example, (H2O2) and (H2O2)3 are formulas.

Given a formula, output the count of all elements as a string in the following form: the first name (in sorted order), followed by its count (if that count is more than 1), followed by the second name (in sorted order), followed by its count (if that count is more than 1), and so on.

Example 1:

Input: 
formula = "H2O"
Output: "H2O"
Explanation: 
The count of elements are {'H': 2, 'O': 1}.



Example 2:

Input: 
formula = "Mg(OH)2"
Output: "H2MgO2"
Explanation: 
The count of elements are {'H': 2, 'Mg': 1, 'O': 2}.



Example 3:

Input: 
formula = "K4(ON(SO3)2)2"
Output: "K4N2O14S4"
Explanation: 
The count of elements are {'K': 4, 'N': 2, 'O': 14, 'S': 4}.



Note:
All atom names consist of lowercase letters, except for the first character which is uppercase.
The length of formula will be in the range [1, 1000].
formula will only consist of letters, digits, and round parentheses, and is a valid formula as defined in the problem.

Source: leetcode | Difficulty: MEDIUM_HARD
Tags: String algorithms

Sorting

Data structures

Starter code: class Solution:
    def countOfAtoms(self, formula: str) -> str:
        

Solutions:
class Solution(object):
	loc = 0
	lastloc = -1
	f = ''

	def getNext(self, formular, locked=False):
		stype = 0
		ret = 0
		if self.loc == self.lastloc:
			return (0, 0)
		i = self.loc
		while i < len(formular):
			if stype in (0, 1) and formular[i].isnumeric():
				ret = int(formular[i]) + ret * 10
				stype = 1
			elif stype == 0 and formular[i].isupper():
				stype = 20
				ret = formular[i]
			elif stype in (20, 2) and formular[i].islower():
				stype = 2
				ret += formular[i]
			elif stype == 0 and formular[i] in '()':
				stype = 3 + '()'.index(formular[i])
				ret = formular[i]
			else:
				break
			i += 1
		if not locked:
			self.lastloc = self.loc
			self.loc = i
		return (stype, ret)

	def countOfAtoms(self, formula):
		stk = []
		cnt = {}
		r = ''
		self.loc = 0
		n = self.getNext(formula)
		while n != (0, 0):
			if n[0] in (2, 3, 20):
				stk.append([n[1], 1])
			elif n[0] == 1:
				stk[-1][1] = n[1]
			elif n[0] == 4:
				time = 1
				i = -1
				if self.getNext(formula, True)[0] == 1:
					time = self.getNext(formula)[1]
				while stk[i][0] != '(':
					stk[i][1] *= time
					i -= 1
				stk[i][0] = '$'
			n = self.getNext(formula)
		while any(stk):
			n = stk.pop()
			if n[0] != '$':
				cnt[n[0]] = cnt.get(n[0], 0) + n[1]
		for i in sorted(cnt.keys()):
			r += '%s%d' % (i, cnt[i]) if cnt[i] > 1 else i
		return r


class Solution:

	def countOfAtoms(self, formula):
		stack = [collections.Counter()]
		i = 0
		length = len(formula)
		ss = ''
		while i < length:
			if formula[i].isupper():
				ss = formula[i]
				i += 1
				while i < length:
					if formula[i].islower():
						ss += formula[i]
						i += 1
					else:
						i -= 1
						break
				num = 0
				i += 1
				while i < length:
					if formula[i].isdigit():
						num = 10 * num + int(formula[i])
						i += 1
					else:
						i -= 1
						break
				num = num if num else 1
				stack[-1][ss] = num if ss not in stack[-1] else stack[-1][ss] + num
			elif formula[i] == '(':
				stack.append(collections.Counter())
			elif formula[i] == ')':
				i += 1
				multiplier = 0
				while i < length:
					if formula[i].isdigit():
						multiplier = 10 * multiplier + int(formula[i])
						i += 1
					else:
						i -= 1
						break
				cur = stack.pop()
				for (n, v) in list(cur.items()):
					stack[-1][n] = stack[-1][n] + v * multiplier if n in stack[-1] else v * multiplier
			i += 1
		return ''.join((name + (str(stack[-1][name]) if stack[-1][name] > 1 else '') for name in sorted(stack[-1])))


class Solution:

	def countOfAtoms(self, formula):
		stack = []
		mapping = collections.defaultdict(int)
		s = ''
		num = 0
		i = 0
		while i < len(formula):
			c = formula[i]
			if c.islower():
				s += c
				i += 1
			elif c.isdigit():
				num = 10 * num + int(c)
				i += 1
			else:
				if num == 0:
					num = 1
				if s != '':
					mapping[s] += num
				s = ''
				num = 0
				i += 1
				if c.isupper():
					s = c
				elif c == '(':
					stack.append(mapping)
					mapping = collections.defaultdict(int)
				elif c == ')':
					while i < len(formula) and formula[i].isdigit():
						num = 10 * num + int(formula[i])
						i += 1
					prev = stack.pop()
					for k in mapping:
						prev[k] += mapping[k] * num
					mapping = prev
					num = 0
		if num == 0:
			num = 1
		if s != '':
			mapping[s] += num
		ans = []
		for (k, v) in sorted(mapping.items()):
			ans.append(k)
			if v > 1:
				ans.append(str(v))
		return ''.join(ans)


class Solution:

	def countOfAtoms(self, formula):
		dict = self.helper(formula)
		result = ''
		for key in sorted(dict):
			result += str(key) + (str(dict[key]) if dict[key] > 1 else '')
		return result

	def helper(self, formula):
		result = {}
		j = 0
		element = ''
		while j < len(formula):
			if formula[j] == '(':
				if len(element) > 0:
					result[element] = (result[element] if element in result else 0) + 1
					element = ''
				count = 1
				end = j + 1
				while end < len(formula) and count > 0:
					if formula[end] == '(':
						count += 1
					if formula[end] == ')':
						count -= 1
					end += 1
				dict = self.helper(formula[j + 1:end - 1])
				i = end
				while i < len(formula) and formula[i].isdigit():
					i += 1
				m = int(formula[end:i])
				for (key, value) in list(dict.items()):
					result[key] = value * m + (result[key] if key in result else 0)
			elif formula[j].isdigit():
				i = j + 1
				while i < len(formula) and formula[i].isdigit():
					i += 1
				result[element] = (result[element] if element in result else 0) + int(formula[j:i])
				element = ''
			else:
				if len(element) > 0:
					result[element] = (result[element] if element in result else 0) + 1
					element = ''
				i = j + 1
				while i < len(formula) and formula[i].islower():
					i += 1
				element = formula[j:i]
			j = i
		if not formula[j - 1].isdigit():
			result[element] = (result[element] if element in result else 0) + 1
		return result


class Solution:

	def countOfAtoms(self, formula):
		formula += ' '
		(stack, i) = ([{}], 0)
		while i < len(formula):
			if formula[i] >= 'A' and formula[i] <= 'Z':
				j = i + 1
				while formula[j] >= 'a' and formula[j] <= 'z':
					j += 1
				k = j
				while formula[k] >= '0' and formula[k] <= '9':
					k += 1
				t = 1 if k == j else int(formula[j:k])
				stack[-1][formula[i:j]] = stack[-1].get(formula[i:j], 0) + t
				i = k
			elif formula[i] == '(':
				stack.append({})
				i += 1
			elif formula[i] == ')':
				j = i + 1
				while formula[j] >= '0' and formula[j] <= '9':
					j += 1
				t = 1 if j == i + 1 else int(formula[i + 1:j])
				for x in stack[-1]:
					stack[-2][x] = stack[-2].get(x, 0) + stack[-1][x] * t
				stack.pop()
				i = j
			else:
				i += 1
		(a, s) = (sorted(stack[-1].keys()), '')
		for x in a:
			s += x if stack[-1][x] == 1 else x + str(stack[-1][x])
		return s


class Solution:

	def __init__(self):
		self.i = 0
		self.formula = ''

	@staticmethod
	def _merge_into(dict_main, dict_added):
		for (key, val) in list(dict_added.items()):
			dict_main[key] = dict_main.get(key, 0) + val

	@staticmethod
	def _multiply(dict_main, mult):
		return {k: v * mult for (k, v) in list(dict_main.items())}

	def cur(self):
		return self.formula[self.i]

	def consume(self):
		res = self.formula[self.i]
		self.i += 1
		return res

	def atend(self):
		return self.i == len(self.formula)

	def countOfAtoms(self, formula):
		self.formula = formula
		res = {}
		while not self.atend():
			self._merge_into(res, self.process_next_part())
		res_str = []
		for (key, value) in sorted(res.items()):
			res_str.append(key)
			if value != 1:
				res_str.append(str(value))
		return ''.join(res_str)

	def process_next_part(self):
		cur = self.cur()
		if cur.isalpha():
			name = self._consume_name()
			value = self._consume_digit()
			return {name: value}
		elif cur == '(':
			local_res = {}
			self.consume()
			while self.cur() != ')':
				self._merge_into(local_res, self.process_next_part())
			self.consume()
			value = self._consume_digit()
			return self._multiply(local_res, value)
		else:
			raise Exception('char unkwon')

	def _consume_name(self):
		name = self.consume()
		if not self.atend() and self.cur().isalpha() and (self.cur().lower() == self.cur()):
			return name + self.consume()
		return name

	def _consume_digit(self):
		if self.atend() or not self.cur().isdigit():
			return 1
		res_str = self.consume()
		while not self.atend() and self.cur().isdigit():
			res_str += self.consume()
		return int(res_str)


class Solution:

	def countOfAtoms(self, formula):

		def merge(d1, d2):
			for (k, v) in list(d2.items()):
				if k in d1:
					d1[k] += v
				else:
					d1[k] = v

		def mul(d, multiplier):
			for k in list(d.keys()):
				d[k] *= multiplier

		def helper(start):
			c = formula[start]
			els = {}
			end = start
			if c == '(':
				end += 1
				while formula[end] != ')':
					(end, els2) = helper(end)
					merge(els, els2)
				end += 1
			elif 'A' <= c <= 'Z':
				end += 1
				while end < len(formula) and 'a' <= formula[end] <= 'z':
					end += 1
				els[formula[start:end]] = 1
			else:
				raise ValueError('Unexpected token')
			end2 = end
			while end2 < len(formula) and '0' <= formula[end2] <= '9':
				end2 += 1
			count = int(formula[end:end2]) if end2 != end else 1
			mul(els, count)
			return (end2, els)
		(i, result) = helper(0)
		while i < len(formula):
			(i, els) = helper(i)
			merge(result, els)
		keys = sorted(result.keys())
		s = ''
		for k in keys:
			if result[k] > 1:
				s += '%s%d' % (k, result[k])
			else:
				s += k
		return s


class Solution:

	def countOfAtoms(self, formula):
		atoms = self.helper(formula)
		return ''.join([e + (str(c) if c > 1 else '') for (e, c) in sorted(atoms.items())])

	def helper(self, formula):
		i = 0
		atoms = {}
		while i < len(formula):
			if formula[i] == '(':
				end = self.end_parens(formula, i)
				atom = self.helper(formula[i + 1:end])
			else:
				end = self.end_element(formula, i)
				atom = {formula[i:end + 1]: 1}
			i = end + 1
			mult = 0
			while i < len(formula) and formula[i].isdigit():
				mult = mult * 10 + int(formula[i])
				i += 1
			if mult == 0:
				mult = 1
			for (element, count) in atom.items():
				atoms.setdefault(element, 0)
				atoms[element] += count * mult
		return atoms

	def end_element(self, formula, i):
		i += 1
		while i < len(formula) and formula[i].islower():
			i += 1
		return i - 1

	def end_parens(self, formula, start):
		parens = 1
		i = start + 1
		while i < len(formula) and parens != 0:
			if formula[i] in '()':
				parens += 1 if formula[i] == '(' else -1
			i += 1
		return i - 1


class Solution:

	def countOfAtoms(self, formula):
		stack = [collections.Counter()]
		i = 0
		while i < len(formula):
			if formula[i] == '(':
				stack.append(collections.Counter())
				print('new stack', stack)
				i += 1
			elif formula[i] == ')':
				i += 1
				num_start = i
				while i < len(formula) and formula[i].isdigit():
					i += 1
				mult = formula[num_start:i]
				if mult is '':
					mult = 1
				else:
					mult = int(formula[num_start:i])
				top = stack[-1]
				print('top: ', top)
				print('new mult: ', mult)
				stack.pop()
				print('poped top stack: ', stack)
				for name in top:
					stack[-1][name] = stack[-1][name] + top[name] * mult
			else:
				name = formula[i]
				i += 1
				while i < len(formula) and formula[i].islower():
					name += formula[i]
					i += 1
				print('name of atom: ', name)
				num_start = i
				while i < len(formula) and formula[i].isdigit():
					i += 1
				mult = formula[num_start:i]
				if mult is '':
					mult = 1
				else:
					mult = int(formula[num_start:i])
				print('mult of {}: '.format(name), mult)
				stack[-1][name] += mult
				print('updated current stack: ', stack)
		return ''.join((name + (str(stack[-1][name]) if stack[-1][name] > 1 else '') for name in sorted(stack[-1])))