2019-1A Problem C

Final solution

// package, import, etc...

func start() {
  // read T, repeat test() T times...
}

type SuffixTree struct {
	strsEndsIn []int
	subTree    map[byte]*SuffixTree
}

func newSuffixTree() *SuffixTree {
	return &SuffixTree{strsEndsIn: make([]int, 0), subTree: make(map[byte]*SuffixTree)}
}

func (t *SuffixTree) addInTree(suffix string, stridx int) {
	leaf := t
	for i := 0; i < len(suffix); i++ {
		char := suffix[i]
		if _, ok := leaf.subTree[char]; !ok {
			leaf.subTree[char] = newSuffixTree()
		}
		leaf = leaf.subTree[char]
	}
	leaf.strsEndsIn = append(leaf.strsEndsIn, stridx)
}

func (t *SuffixTree) find(suffix string) []int {
	leaf := t
	for i := 0; i < len(suffix); i++ {
		char := suffix[i]
		if _, ok := leaf.subTree[char]; !ok {
			return nil
		}
		leaf = leaf.subTree[char]
	}
	return leaf.strsEndsIn
}

func test() {
	var N int
	mustReadLineOfInts(&N)
	words := make([]string, 0, N)
	// t := newSuffixTree()
	maxLength := 0
	for i := 0; i < N; i++ {
		word := mustReadLine()
		words = append(words, word)
		if maxLength < len(word) {
			maxLength = len(word)
		}
		// stridx := len(words)
		// for i := 0; i < len(word); i++ {
		// 	t.addInTree(word[i:], stridx)
		// }
	}
	maxSuffixLength := make([]int, N)
	for i := 0; i < N; i++ {
		maxSuffixLength[i] = len(words[i])
	}
	pairedWordsCount := 0
	// for thisWord := 0; thisWord < N; thisWord++ {
	// 	if !wordsRemaining[thisWord] {
	// 		continue
	// 	}
	// 	word := words[thisWord]
	// 	for suf := 0; suf < len(word); suf++ {
	// 		_strs := t.find(word[suf:])
	// 		strsWithThisSuf := make([]int, 0, len(_strs))
	// 		for _, stridx := range _strs {
	// 			if wordsRemaining[stridx] {
	// 				strsWithThisSuf = append(strsWithThisSuf, stridx)
	// 				if len(strsWithThisSuf) > 2 {
	// 					break
	// 				}
	// 			}
	// 		}
	// 		if len(strsWithThisSuf) > 2 {
	// 			wordsRemaining[thisWord] = false
	// 			break
	// 		}
	// 		if len(strsWithThisSuf) == 2 {
	// 			otherWord := strsWithThisSuf[0]
	// 			if otherWord == thisWord {
	// 				otherWord = strsWithThisSuf[1]
	// 			}
	// 			pairedWordsCount += 2
	// 			wordsRemaining[otherWord] = false
	// 			wordsRemaining[thisWord] = false
	// 			break
	// 		}
	// 		assert(len(strsWithThisSuf) == 1)
	// 	}
	// }

	for suffixLength := maxLength; suffixLength >= 1; suffixLength-- {
		// debug(fmt.Sprintf("suffixLength = %v", suffixLength))
		t := newSuffixTree()
		for thisWord := 0; thisWord < len(words); thisWord++ {
			if maxSuffixLength[thisWord] < suffixLength {
				continue
			}
			word := words[thisWord]
			if len(word) < suffixLength {
				continue
			}
			thisSuffix := word[len(word)-suffixLength:]
			// debug(fmt.Sprintf("Adding %v", thisSuffix))
			t.addInTree(thisSuffix, thisWord)
		}
		for thisWord := 0; thisWord < len(words); thisWord++ {
			if maxSuffixLength[thisWord] < suffixLength {
				continue
			}
			word := words[thisWord]
			if len(word) < suffixLength {
				continue
			}
			thisSuffix := word[len(word)-suffixLength:]
			// debug(thisSuffix)
			_foundWords := t.find(thisSuffix)
			foundWords := make([]int, 0, len(_foundWords))
			for _, w := range _foundWords {
				if maxSuffixLength[w] >= suffixLength {
					foundWords = append(foundWords, w)
				}
			}
			if len(foundWords) >= 2 {
				otherWord := foundWords[0]
				if otherWord == thisWord {
					otherWord = foundWords[1]
				}
				maxSuffixLength[thisWord] = 0
				maxSuffixLength[otherWord] = 0
				pairedWordsCount += 2
				// debug(fmt.Sprintf("Pair: %v and %v. Found %v other words", word, words[otherWord], len(foundWords)-1))
				for i := 2; i < len(foundWords); i++ {
					maxSuffixLength[foundWords[i]] = suffixLength - 1
					// debug(fmt.Sprintf("Decreasing suffixLength for %v to %v", words[foundWords[i]], maxSuffixLength[foundWords[i]]))
				}
				continue
			}
		}
	}
	fmt.Fprintf(stdout, "%d\n", pairedWordsCount)
}

// boilerplate omitted...

Sample test cases given in question

4
2
TARPOL
PROL
3
TARPOR
PROL
TARPRO
6
CODEJAM
JAM
HAM
NALAM
HUM
NOLOM
4
PI
HI
WI
FI

Other files

1
8
BBEJAM
CODEJAM
AAEJAM
BBBJAM
EAJAM
CCCJAM
BEJAM
DAJAM