LeetCode/go/number-of-atoms.go

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
}