leetcode-go/leetcode/medium/640.solve-the-equation.go

func solveEquation(equation string) string {
	part := strings.Split(equation, "=")
	x1, a1 := parse(part[0])
	x2, a2 := parse(part[1])
	if x1 == x2 && a1 == a2 {
		return "Infinite solutions"
	} else if x1 != x2 {
		return fmt.Sprintf("x=%d", (a2-a1)/(x1-x2))
	} else {
		return "No solution"
	}
}

func parse(eq string) (x int, a int) {
	pre, n := 0, len(eq)
	for i := range eq {
		if eq[i] == 'x' || i == n-1 || eq[i+1] == '-' || eq[i+1] == '+' {
			if eq[i] == 'x' {
				switch eq[pre:i] {
				case "", "+": // Handle special case
					x += 1
				case "-":
					x -= 1
				default:
					num, _ := strconv.Atoi(eq[pre:i])
					x += num
				}
			} else {
				num, _ := strconv.Atoi(eq[pre : i+1])
				a += num
			}
			pre = i + 1
		}
	}
	return
}