aboutsummaryrefslogtreecommitdiff
path: root/number_atoms.py
diff options
context:
space:
mode:
Diffstat (limited to 'number_atoms.py')
-rw-r--r--number_atoms.py47
1 files changed, 47 insertions, 0 deletions
diff --git a/number_atoms.py b/number_atoms.py
new file mode 100644
index 0000000..86fa2d7
--- /dev/null
+++ b/number_atoms.py
@@ -0,0 +1,47 @@
+# 726. Number of Atoms
+from collections import defaultdict
+
+"""
+given a string 'formula' representing a chemical formula, return the count of
+each atom. the atomic element always stats with an uppercase character, then
+zer or more lowercase letters, representing the name. one or more digits
+representing that element's count may follow if the count is greater than 1.
+if the count is 1, no digits will follow. return the count of all elements as
+a string in the following form: the first name, followed by its count,
+followed by the second name, followed by its count, and so on.
+"""
+
+
+class Solution(object):
+ def countOfAtoms(self, formula):
+ """
+ :type formula: str
+ :rtype: str
+ """
+ dic, coeff, stack, elem, cnt, i = defaultdict(int), 1, [], "", 0, 0
+ for c in formula[::-1]:
+ if c.isdigit():
+ cnt += int(c) * (10**i)
+ i += 1
+ elif c == ")":
+ stack.append(cnt)
+ coeff *= cnt
+ i = cnt = 0
+ elif c == "(":
+ coeff //= stack.pop()
+ i = cnt = 0
+ elif c.isupper():
+ elem += c
+ dic[elem[::-1]] += (cnt or 1) * coeff
+ elem = ""
+ i = cnt = 0
+ elif c.islower():
+ elem += c
+ return "".join(k + str(v > 1 and v or "") for k, v in sorted(dic.items()))
+
+
+if __name__ == "__main__":
+ obj = Solution()
+ print(obj.countOfAtoms(formula="H2O"))
+ print(obj.countOfAtoms(formula="Mg(OH)2"))
+ print(obj.countOfAtoms(formula="K4(ON(SO3)2)2"))