競プロ落とし穴 #1 SetでDFS

競プロ落とし穴シリーズ

シリーズと書いてるけど多分これが初回で最終回
競プロの問題を解いていてつまづいたところをメモしていくよ

今回解いていた問題

出現する文字列(たかだか10種類)について全探索を行って条件を満たす解を探す問題

以下では簡略化したコードで考えていくよ

間違ったコード

func AssignWithSet(N int) {
	set := make(map[int]struct{})
	for i := 0; i < N; i++ {
		set[i] = struct{}{}
	}
	ret := make(map[int]int)
	used := make(map[int]bool)
	for v := range set {
		ret[v] = -1
	}
	scnt := 0
	var DFSWithSet func(int)
	DFSWithSet = func(num int) {
		if num == len(set) {
			scnt++
			fmt.Printf("%v\n", ret)
			return
		}
		for v := range set {
			if ret[v] != -1 {
				continue
			}
			for i := 0; i < 10; i++ {
				if used[i] {
					continue
				}
				used[i] = true
				ret[v] = i
				DFSWithSet(num + 1)
				used[i] = false
				ret[v] = -1
			}
		}
	}
	DFSWithSet(0)
	fmt.Printf("total: %d\n", scnt)
}

要素数NのSetに対してそれぞれ重複なく0~9の数字を割り当てている
でも実はかなり無駄がある

変数scntはすべての要素に数字を割り当てた状況を何回作ったかを数えてくれてる
実際にプログラムを実行してNとscntの関係を見てみると

N	1	2	3	4
scnt	10	180	4320	120960

正しいコード

func AssignWithArray(N int) {
	arr := make([]int, 0)
	for i := 0; i < N; i++ {
		arr = append(arr, i)
	}
	ret := make(map[int]int)
	used := make(map[int]bool)
	for v := range arr {
		ret[v] = -1
	}
	acnt := 0
	var DFSWithArray func(int)
	DFSWithArray = func(num int) {
		if num == len(arr) {
			acnt++
			fmt.Printf("%v\n", ret)
			return
		}
		v := arr[num]
		for i := 0; i < 10; i++ {
			if used[i] {
				continue
			}
			used[i] = true
			ret[v] = i
			DFSWithArray(num + 1)
			used[i] = false
		}
	}
	DFSWithArray(0)
	fmt.Printf("total: %d\n", acnt)
}

さっきのコードと同じことを今度は配列(厳密にはSlice)で実装してみた
こっちは一つの割当をちょうど1回ずつしかみない無駄のない挙動をしてくれる

変数acntはすべての要素に数字を割り当てた状況を何回作ったかを数えてくれてる
同様にNとacntの関係を見てみると

N	1	2	3	4
acnt	10	90	720	5040

挙動の違い

今度はscntとacntを比べてみると

N	1	2	3	4
scnt	10	180	4320	120960
acnt	10	90	720	5040
scnt/acnt	1	2	6	24

ちょうどN!倍だけscntの方が多く探索しているのがわかる

AssignWithSetでは添字ではなく、すでに割り当てた要素の数で管理しているから探索後に割当を取り消してしまう。この分の差がN!倍につながっている

要はSetをArrayに変換するのを怠けて実装してしまったから探索の無駄が生まれてTLEしてしまったというのが今回の躓いたポイント！
(SetでもうまくやればDFSできるかも？)

最後にACとなったコードも貼っておくよ

package main

import (
	"bufio"
	"fmt"
	"math"
	"os"
)

const (
	INF = int(1 << 61)
	MOD = int(1e9 + 7)
	// MOD = 998244353
)

func main() {
	defer _w.Flush()
	S := make([]string, 3)
	for i := 0; i < 3; i++ {
		fmt.Fscan(_r, &S[i])
	}
	Solve(S)
}

func Solve(S []string) {
	set := make(map[rune]struct{})
	for i := 0; i < 3; i++ {
		for _, c := range S[i] {
			set[c] = struct{}{}
		}
	}
	if len(set) > 10 {
		fmt.Fprintln(_w, "UNSOLVABLE")
		return
	}
	arr := make([]rune, 0)
	for k, _ := range set {
		arr = append(arr, k)
	}
	mp := make(map[rune]int)
	used := make(map[int]bool)
	for _, c := range arr {
		mp[c] = -1
	}
	convert := func(T string) int {
		ret := 0
		for i := 0; i < len(T); i++ {
			ret *= 10
			ret += mp[rune(T[i])]
		}
		return ret
	}
	check := func() bool {
		for i := 0; i < 3; i++ {
			if mp[rune(S[i][0])] == 0 {
				return false
			}
		}
		return convert(S[0])+convert(S[1]) == convert(S[2])
	}
	var dfs func(int) bool
	dfs = func(num int) bool {
		if num == len(arr) {
			if check() {
				for i := 0; i < 3; i++ {
					fmt.Printf("%d\n", convert(S[i]))
				}
				return true
			} else {
				return false
			}
		}
		c := arr[num]
		for i := 0; i < 10; i++ {
			if used[i] {
				continue
			}
			mp[c] = i
			used[i] = true
			if dfs(num + 1) {
				return true
			}
			mp[c] = -1
			used[i] = false
		}
		return false
	}
	if !dfs(0) {
		fmt.Fprintln(_w, "UNSOLVABLE")
	}
}

var _r, _w = bufio.NewReader(os.Stdin), bufio.NewWriter(os.Stdout)