1989 solution - version #2

digits2.go

package main

import (
	"fmt"
	"log"
	"math"
	"os"
	"sort"
	"strconv"
	"strings"
)

// Lookup tables for fast factorial and integer square root
var factLookup, sqrtLookup map[int64]int64

const (
	maxFactorial = 20     // Pre-calculate n! up to this
	maxSqrt      = 100000 // Pre-calculate sqrt[n] up to this
	maxSqrt2     = maxSqrt * maxSqrt
)

func init() {
	factLookup = make(map[int64]int64)
	sqrtLookup = make(map[int64]int64)
	factLookup[0] = 0
	var i, fact int64
	fact = 2
	// We are starting from 3, since 1! = 1 and 2! = 2
	for i = 3; i <= maxFactorial; i++ {
		fact *= i
		factLookup[i] = fact
	}
	for i = 1; i <= maxSqrt; i++ {
		sqrtLookup[i*i] = i
	}
}

// fact calculates n! using lookup table, and returns -1 for invalid inputs
func fact(n int64) int64 {
	if n < 0 || n > maxFactorial || (n < 3 && n != 0) {
		return -1
	}
	return factLookup[n]
}

// sqrt calculates sqrt[n] using lookup table, and returns -1 for invalid inputs
// or non-integer square roots
func sqrt(n int64) int64 {
	if n < 2 || n > maxSqrt2 {
		return -1
	}
	if r, ok := sqrtLookup[n]; ok {
		return r
	}
	return -1
}

const InvalidPow = 9223372036854775807 // max int64

// pow returns a^b, if both of them are small enough, or InvalidPow if the result is invalid
// FIXME: Once we support ratios, we should support a^r where r is a ratio, too.
func pow(a, b int64) int64 {
	if a == 0 || b <= 0 {
		return InvalidPow
	} else if a == 0 {
		return 0
	} else if b == 0 {
		return 1
	}
	if b > 15 && (a > 15 || a < -15) {
		return InvalidPow
	}
	if p := math.Pow(float64(a), float64(b)); math.Abs(p) > InvalidPow {
		// Overflow
		return InvalidPow
	} else {
		return int64(p)
	}
}

type Op byte // Operators

const (
	OpNull Op = iota
	// Binary ops start here
	OpAdd
	OpSub
	OpMul
	OpDiv
	OpPow
	// Unary ops start here
	OpFact  // factorial
	OpSqrt  // square root
	OpMinus // unary minus
)

var opNames = map[Op]string{
	OpNull:  "NULL",
	OpAdd:   "+",
	OpSub:   "-",
	OpMul:   "*",
	OpDiv:   "/",
	OpPow:   "^",
	OpFact:  "!",
	OpSqrt:  "sqrt",
	OpMinus: "-",
}

type Node struct {
	left, right *Node
	val         int64 // Number stored in this node

	// Operation to be perfomed on this node. Require left != nil && right != nil for binary
	// and left != nil && right = nil for unary operators.
	op Op
}

// String returns a formula for n, sometimes (always *sigh*) with excessive paranthesis
func (n *Node) String() string {
	var left, right string
	if n.left != nil {
		left = n.left.String()
	}
	if n.right != nil {
		right = n.right.String()
	}
	if left == "" && right == "" {
		return strconv.FormatInt(n.val, 10)
	} else {
		switch n.op {
		case OpAdd:
			return fmt.Sprintf("%s + (%s)", left, right)
		case OpSub:
			return fmt.Sprintf("%s - (%s)", left, right)
		case OpMul, OpDiv, OpPow:
			return fmt.Sprintf("(%s) %s (%s)", left, opNames[n.op], right)
		case OpFact:
			return fmt.Sprintf("(%s)!", left)
		case OpSqrt, OpMinus:
			return fmt.Sprintf("%s(%s)", opNames[n.op], left)
		default:
			return "<UNDEFINED>" // Unreachable
		}
	}
}

func (n *Node) Depth() int64 {
	if n.left == nil && n.right == nil {
		return 0
	}
	depth := n.left.Depth()
	if n.right != nil {
		if d := n.right.Depth(); d > depth {
			depth = d
		}
	}
	return depth + 1
}

func (n *Node) Equal(n1 *Node) bool {
	if n.val != n1.val || n.op != n1.op {
		return false
	}
	if (n.left == nil && n1.left != nil) ||
		(n.left != nil && n1.left == nil) ||
		(n.left != nil && n1.left != nil && !n.left.Equal(n1.left)) {
		return false
	}
	if (n.right == nil && n1.right != nil) ||
		(n.right != nil && n1.right == nil) ||
		(n.right != nil && n1.right != nil && !n.right.Equal(n1.right)) {
		return false
	}
	return true
}

// Solution for the formula.
// start and end indicates that we used digits[start:end] for this formula,
// where digits is the original digits string. We cannot use binary operator for s1, s2
// if s1.end != s2.start
type Solution struct {
	val        int64
	start, end int
}

var NoSolution Solution

func init() {
	NoSolution = Solution{val: 0, start: -1, end: -1}
}

// Global variables are bad for you health
var solutions map[Solution][]*Node // solutions found so far
var maxDepth int64                 // If positive, only search for formulas of up to this level. If zero, only stores the first solution.

func init() {
	solutions = make(map[Solution][]*Node) // I love go's ability to have a struct as a key
}

// Add adds a new formula for s, but only if it's unique and has reasonable depth.
// If v == nil, seed solutions with initial digits.
func (s Solution) Add(v *Node) {
	if maxDepth == 0 && solutions[s] != nil {
		// Do nothing, we only need the first solution
		return
	}
	if v == nil {
		v = &Node{val: s.val}
	}
	if maxDepth != 0 && v.Depth() > maxDepth {
		return
	}
	for _, v1 := range solutions[s] {
		if v.Equal(v1) {
			return
		}
	}
	solutions[s] = append(solutions[s], v)
}

// Apply an unary operator to this solution, if possible, and add to all solutions
// found so far. Returns a list of solutions that can be received this way,
// including the original (= do nothing).
func (s Solution) Unary(op Op) Solution {
	if s.val == 0 {
		return s
	}
	var s1 Solution
	switch op {
	case OpMinus:
		s1 = Solution{val: -s.val, start: s.start, end: s.end}
	case OpFact:
		if f := fact(s.val); f != -1 {
			s1 = Solution{val: f, start: s.start, end: s.end}
		} else {
			return NoSolution
		}
	case OpSqrt:
		if sq := sqrt(s.val); sq != -1 {
			s1 = Solution{val: sq, start: s.start, end: s.end}
		} else {
			return NoSolution
		}
	default:
		return NoSolution
	}
	for _, n := range solutions[s] {
		if n.op == OpMinus && op == OpMinus { // Do not apply to unary minus in a row
			continue
		}
		if n.op == OpPow && op == OpSqrt { // we only want to allow sqrt(sqrt(x^4)), not sqrt(sqrt(x)^4) and all other variations
			continue
		}
		s1.Add(&Node{op: op, left: n})
	}
	return s1
}

// Apply a binary operator to this two solutions, if possible, and add to all solutions
// found so far. Returns a list of solutions that can be received this way.
func (s1 Solution) Binary(op Op, s2 Solution) Solution {
	if s1.end != s2.start {
		return NoSolution
	}
	var s3 Solution
	switch op {
	case OpAdd:
		s3 = Solution{
			val:   s1.val + s2.val,
			start: s1.start, end: s2.end,
		}
	case OpSub:
		s3 = Solution{
			val:   s1.val - s2.val,
			start: s1.start, end: s2.end,
		}
	case OpMul:
		s3 = Solution{
			val:   s1.val * s2.val,
			start: s1.start, end: s2.end,
		}
	case OpDiv:
		if s2.val == 0 {
			return NoSolution
		}
		if r := s1.val / s2.val; r*s2.val == s1.val {
			s3 = Solution{
				val:   r,
				start: s1.start, end: s2.end,
			}
		} else {
			return NoSolution
		}
	case OpPow:
		if p := pow(s1.val, s2.val); p != InvalidPow {
			s3 = Solution{
				val:   p,
				start: s1.start, end: s2.end,
			}
		} else {
			return NoSolution
		}
	default:
		return NoSolution
	}
	for _, n1 := range solutions[s1] {
		for _, n2 := range solutions[s2] {
			if op == OpMinus && n2.op == OpMinus { // do not generate x - (-y)
				continue
			}
			s3.Add(&Node{
				op:    op,
				left:  n1,
				right: n2,
			})
		}
	}
	return s3
}

// AllUnary generates all possible solutions we can get from s using unary operations,
// including itself (= no operation was applied).
func (s Solution) AllUnary() []Solution {
	if s.val == 0 {
		return []Solution{s}
	}
	result := []Solution{s}
	s1 := s.Unary(OpMinus)
	if s1 != NoSolution {
		result = append(result, s1)
	}
	if s.val < 0 {
		if s1 == NoSolution {
			return result // we don't want to apply fact, sqrt to a negative number
		} else {
			s = s1
		}
	}
	// Factorial - it is never a perfect square
	for f := s.Unary(OpFact); f != NoSolution; f = f.Unary(OpFact) {
		result = append(result, f)
		result = append(result, f.Unary(OpMinus)) // And -s! was added to solution pool at this point
	}
	// Square root
	for sq := s.Unary(OpSqrt); sq != NoSolution; sq = sq.Unary(OpSqrt) {
		result = append(result, sq)
		for f := sq.Unary(OpFact); f != NoSolution; f = f.Unary(OpFact) {
			result = append(result, f)
		}
	}
	return uniq(result)
}

// AllBinary generates all possible binary solutions for s1 and s2.
func (s1 Solution) AllBinary(s2 Solution) []Solution {
	result := []Solution{}
	for op := OpAdd; op <= OpPow; op++ {
		for _, s3 := range s1.AllUnary() {
			for _, s4 := range s2.AllUnary() {
				if s5 := s3.Binary(op, s4); s5 != NoSolution {
					result = append(result, s5)
				}
			}
		}
	}
	return uniq(result)
}

// MakeAllTernary does what it name implies :)
func MakeAllTernary(s1, s2, s3 Solution) []Solution {
	result := []Solution{}
	for _, s12 := range s1.AllBinary(s2) {
		result = append(result, s12.AllBinary(s3)...)
	}
	for _, s23 := range s2.AllBinary(s3) {
		result = append(result, s1.AllBinary(s23)...)
	}
	return uniq(result)
}

// MakeAllQuaternary does what it name implies :)
func MakeAllQuaternary(s1, s2, s3, s4 Solution) []Solution {
	result := []Solution{}
	for _, s123 := range MakeAllTernary(s1, s2, s3) {
		result = append(result, s123.AllBinary(s4)...)
	}
	for _, s234 := range MakeAllTernary(s2, s3, s4) {
		result = append(result, s1.AllBinary(s234)...)
	}
	for _, s12 := range s1.AllBinary(s2) {
		for _, s34 := range s3.AllBinary(s4) {
			result = append(result, s12.AllBinary(s34)...)
		}
	}
	return uniq(result)
}

// uniq returns only unique solutions from the list
func uniq(l []Solution) []Solution {
	m := make(map[Solution]bool)
	for _, n := range l {
		m[n] = true
	}
	var result []Solution
	for k := range m {
		result = append(result, k)
	}
	return result
}

func atos(a string, start, end int) Solution {
	n, err := strconv.Atoi(a[start:end])
	if err != nil {
		log.Fatalf("Cannot convert %s to number\n", a)
	}
	s := Solution{val: int64(n), start: start, end: end}
	s.Add(nil)
	return s
}

func atoi(s string) int64 {
	n, err := strconv.Atoi(s)
	if err != nil {
		log.Fatalf("Cannot convert %s to number\n", s)
	}
	return int64(n)
}

// SolutionSlice provides sort.Sortable interface for []Solution, ignoring level
type SolutionSlice []Solution

func (p SolutionSlice) Len() int {
	return len(p)
}

func (p SolutionSlice) Less(i, j int) bool {
	return p[i].val < p[j].val
}

func (p SolutionSlice) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}

func sorted(n []Solution) []Solution {
	a := SolutionSlice(n)
	sort.Sort(a)
	return ([]Solution)(a)
}

func generateFormulas(digits string, min, max int64) []Solution {
	if _, err := strconv.Atoi(digits); len(digits) != 4 || len(digits) == 0 || err != nil {
		log.Fatalf(`Please run me like this: 'digits2 abcd min max', where abcd is a string of 4 digits.
I'll print all formulas for numbers between min and max.
An optional fourth argument defines depth of the search and prints all the solutions.
If missing, it will only find a first solution for each number.`)
	}
	n1, n2, n3, n4 := atos(digits, 0, 1), atos(digits, 1, 2), atos(digits, 2, 3), atos(digits, 3, 4)
	n12, n23, n34 := atos(digits, 0, 2), atos(digits, 1, 3), atos(digits, 2, 4)
	n123, n234 := atos(digits, 0, 3), atos(digits, 1, 4)
	n1234 := atos(digits, 0, 4)

	formulas := []Solution{}

	result := MakeAllQuaternary(n1, n2, n3, n4)
	result = append(result, MakeAllTernary(n12, n2, n3)...)
	result = append(result, MakeAllTernary(n1, n23, n4)...)
	result = append(result, MakeAllTernary(n1, n2, n34)...)
	result = append(result, n123.AllBinary(n4)...)
	result = append(result, n12.AllBinary(n34)...)
	result = append(result, n1.AllBinary(n234)...)
	result = append(result, n1234.AllUnary()...)
	for _, n := range uniq(result) {
		if n.val >= min && n.val <= max {
			formulas = append(formulas, n)
		}
	}
	return sorted(formulas)
}

// func main() {
// 	if len(os.Args) == 2 {
// 		fmt.Printf("Printing all unary\n")
// 		for _, f := range atos(os.Args[1]).AllUnary() {
// 			for _, n := range solutions[f] {
// 				fmt.Printf("%d\t= [%d] %s\n", f.val, n.Depth(), n)
// 			}
// 		}
// 	} else if len(os.Args) == 3 {
// 		fmt.Printf("Printing all binary\n")
// 		for _, f := range atos(os.Args[1]).AllBinary(atos(os.Args[2])) {
// 			for _, n := range solutions[f] {
// 				fmt.Printf("%d\t= [%d] %s\n", f.val, n.Depth(), n)
// 			}
// 		}
// 	} else if len(os.Args) == 4 {
// 		fmt.Printf("Printing all ternarary\n")
// 		for _, f := range MakeAllTernary(atos(os.Args[1]), atos(os.Args[2]), atos(os.Args[3])) {
// 			for _, n := range solutions[f] {
// 				fmt.Printf("%d\t= [%d] %s\n", f.val, n.Depth(), n)
// 			}
// 		}
// 	} else if len(os.Args) == 5 {
// 		fmt.Printf("Printing all quaternary\n")
// 		for _, f := range MakeAllQuaternary(atos(os.Args[1]), atos(os.Args[2]), atos(os.Args[3]), atos(os.Args[4])) {
// 			for _, n := range solutions[f] {
// 				fmt.Printf("%d\t= [%d] %s\n", f.val, n.Depth(), n)
// 			}
// 		}
// 	}
// }

func main() {
	var printAll bool
	if len(os.Args) > 4 {
		maxDepth = atoi(os.Args[4])
		printAll = true
	}
	formulas := generateFormulas(os.Args[1], atoi(os.Args[2]), atoi(os.Args[3]))
	for _, f := range formulas {
		if printAll {
			fmt.Printf("---\nAll formulas for number %d up to depth = %d:\n", f.val, maxDepth)
		} else {
			fmt.Printf("%d\t= ", f.val)
		}
		answer := []string{}
		for _, n := range solutions[f] {
			answer = append(answer, fmt.Sprintf("[%d] %s", n.Depth(), n))
		}
		sort.Strings(answer)
		fmt.Printf("%s\n", strings.Join(answer, "\n"))
	}
}