go: add «726. Number of Atoms»

Signed-off-by: Matej Focko <me@mfocko.xyz>

go/number-of-atoms.go

@@ -0,0 +1,92 @@
+package main
+
+import (
+	"fmt"
+	"slices"
+	"strconv"
+	"unicode"
+)
+
+func countOfAtoms(formula string) string {
+	var stack []map[string]int
+	stack = append(stack, make(map[string]int))
+
+	i := 0
+	for i < len(formula) {
+		if formula[i] == '(' {
+			stack = append(stack, make(map[string]int))
+			i++
+		} else if formula[i] == ')' {
+			count := stack[len(stack)-1]
+			stack = stack[:len(stack)-1]
+
+			i++
+			j := i
+			for j < len(formula) && unicode.IsDigit(rune(formula[j])) {
+				j++
+			}
+
+			multiplier := 1
+			if i < j {
+				parsedMultiplier, err := strconv.ParseInt(formula[i:j], 10, 0)
+				if err != nil {
+					panic(err)
+				}
+
+				multiplier = int(parsedMultiplier)
+			}
+			i = j
+
+			parent := stack[len(stack)-1]
+			for element, count := range count {
+				parent[element] += count * multiplier
+			}
+		} else {
+			j := i + 1
+			for j < len(formula) && unicode.IsLower(rune(formula[j])) {
+				j++
+			}
+
+			element := formula[i:j]
+			i = j
+
+			for j < len(formula) && unicode.IsDigit(rune(formula[j])) {
+				j++
+			}
+
+			count := 1
+			if i < j {
+				parsedCount, err := strconv.ParseInt(formula[i:j], 10, 0)
+				if err != nil {
+					panic(err)
+				}
+
+				count = int(parsedCount)
+			}
+			i = j
+
+			parent := stack[len(stack)-1]
+			parent[element] += count
+		}
+	}
+
+	finalCount := stack[0]
+
+	elements := make([]string, len(finalCount))
+	for element := range finalCount {
+		elements = append(elements, element)
+	}
+	slices.Sort(elements)
+
+	answer := ""
+	for _, element := range elements {
+		answer += element
+
+		count := finalCount[element]
+		if count > 1 {
+			answer += fmt.Sprintf("%d", count)
+		}
+	}
+
+	return answer
+}

go/number-of-atoms_test.go

@@ -0,0 +1,28 @@
+package main
+
+import "testing"
+
+func _test_NumberOfAtoms(t *testing.T, input, expected string) {
+	result := countOfAtoms(input)
+	if result != expected {
+		t.Errorf("Expected ‹%s› for ‹%s›, but got: ‹%s›!", expected, input, result)
+	}
+}
+
+func Test_NumberOfAtoms_Easy(t *testing.T) {
+	_test_NumberOfAtoms(t, "H2", "H2")
+	_test_NumberOfAtoms(t, "O2", "O2")
+	_test_NumberOfAtoms(t, "HO", "HO")
+	_test_NumberOfAtoms(t, "H2O", "H2O")
+	_test_NumberOfAtoms(t, "H2SO4", "H2O4S")
+}
+
+func Test_NumberOfAtoms_Medium(t *testing.T) {
+	_test_NumberOfAtoms(t, "Mg(OH)2", "H2MgO2")
+	_test_NumberOfAtoms(t, "C6H5OH", "C6H6O")
+	_test_NumberOfAtoms(t, "C6H5CH3", "C7H8")
+}
+
+func Test_NumberOfAtoms_Hard(t *testing.T) {
+	_test_NumberOfAtoms(t, "K4(ON(SO3)2)2", "K4N2O14S4")
+}