package alphametics
import (
	"reflect"
	"testing"
)
func TestSolve(t *testing.T) {
	for _, tc := range testCases {
		t.Run(tc.description, func(t *testing.T) {
			s, err := Solve(tc.input)
			switch {
			case tc.errorExpected:
				if err == nil {
					t.Fatalf("Solve(%q) expected error, got: %#v", tc.input, s)
				}
			case err != nil:
				t.Fatalf("Solve(%q)\nexpected: %#v\ngot error: %q", tc.input, tc.expected, err)
			case !reflect.DeepEqual(s, tc.expected):
				t.Fatalf("Solve(%q)\ngot: %#v\nwant:%#v", tc.input, s, tc.expected)
			}
		})
	}
}
func BenchmarkSolve(b *testing.B) {
	if testing.Short() {
		b.Skip("skipping benchmark in short mode.")
	}
	for i := 0; i < b.N; i++ {
		for _, tc := range testCases {
			Solve(tc.input)
		}
	}
}

package alphametics
import (
	"fmt"
	"math"
	"strings"
)
type LetterPlaces map[rune][]int
type FirstMap map[byte]struct{}
type Puzzle struct {
	addLetters   LetterPlaces   // letter places in add tags
	sumLetters   LetterPlaces   // letter places in sum
	firstLetters FirstMap       // first letters, they can not be 0
	letters      []string       // fixed position list of the letters
	pool         []bool         // marks for used/available numbers
	solution     map[string]int // solution in expected form
}
func Solve(puzzle string) (map[string]int, error) {
	pzl := Puzzle{
		addLetters:   LetterPlaces{},
		sumLetters:   LetterPlaces{},
		firstLetters: FirstMap{},
		letters:      []string{},
		pool:         make([]bool, 10), // there are only 10 number to choose from
		solution:     map[string]int{},
	}
	// collect the position of each letter for sum and addition tags
	parts := strings.Split(puzzle, " == ")
	pzl.fill(pzl.sumLetters, parts[1])
	adds := strings.Split(parts[0], " + ")
	for _, a := range adds {
		pzl.fill(pzl.addLetters, a)
	}
	// fill in the letters, order is not important
	for k := range pzl.solution {
		pzl.letters = append(pzl.letters, k)
	}
	// recursive brute-force solver
	if pzl.recurse(0, puzzle) {
		return pzl.solution, nil
	}
	return pzl.solution, fmt.Errorf("no solution")
}
func (pzl Puzzle) fill(lp LetterPlaces, in string) {
	pzl.firstLetters[in[0]] = struct{}{}
	for idx, ch := range in {
		lp[ch] = append(lp[ch], len(in)-idx-1) // reverse the index for proper place-value
		pzl.solution[string(ch)] = 0           // add chars, initial value does not matter
	}
}
func (pzl Puzzle) recurse(pos int, puzzle string) bool {
	if pos == len(pzl.solution) {
		if pzl.checkMap() {
			return true
		}
		return false
	}
	for i, val := range pzl.pool {
		if !val {
			// handle first characters, they can not be 0
			if i == 0 {
				if _, ok := pzl.firstLetters[pzl.letters[pos][0]]; ok {
					continue
				}
			}
			// mark i-th item of the pool used
			pzl.pool[i] = true
			// set the value of the current letter
			pzl.solution[pzl.letters[pos]] = i
			if !pzl.recurse(pos+1, puzzle) {
				// if solution is not right, unmark i-th number in the pool
				pzl.pool[i] = false
				continue
			}
			// right solution
			return true
		}
	}
	return false
}
func (pzl Puzzle) checkMap() bool {
	sum := 0
	for ch, positions := range pzl.sumLetters {
		for _, p := range positions {
			sum += pzl.solution[string(ch)] * int(math.Pow10(p))
		}
	}
	num := 0
	for ch, positions := range pzl.addLetters {
		for _, p := range positions {
			num += pzl.solution[string(ch)] * int(math.Pow10(p))
		}
		// return when solution is doomed...
		if num > sum {
			return false
		}
	}
	return sum == num
}