3045-CountPrefixAndSuffixPairsII.go

package main

// 3045. Count Prefix and Suffix Pairs II
// You are given a 0-indexed string array words.

// Let's define a boolean function isPrefixAndSuffix that takes two strings, str1 and str2:
//     isPrefixAndSuffix(str1, str2) returns true if str1 is both a prefix and a suffix of str2, and false otherwise.

// For example, isPrefixAndSuffix("aba", "ababa") is true because "aba" is a prefix of "ababa" 
// and also a suffix, but isPrefixAndSuffix("abc", "abcd") is false.

// Return an integer denoting the number of index pairs (i, j) 
// such that i < j, and isPrefixAndSuffix(words[i], words[j]) is true.

// Example 1:
// Input: words = ["a","aba","ababa","aa"]
// Output: 4
// Explanation: In this example, the counted index pairs are:
// i = 0 and j = 1 because isPrefixAndSuffix("a", "aba") is true.
// i = 0 and j = 2 because isPrefixAndSuffix("a", "ababa") is true.
// i = 0 and j = 3 because isPrefixAndSuffix("a", "aa") is true.
// i = 1 and j = 2 because isPrefixAndSuffix("aba", "ababa") is true.
// Therefore, the answer is 4.

// Example 2:
// Input: words = ["pa","papa","ma","mama"]
// Output: 2
// Explanation: In this example, the counted index pairs are:
// i = 0 and j = 1 because isPrefixAndSuffix("pa", "papa") is true.
// i = 2 and j = 3 because isPrefixAndSuffix("ma", "mama") is true.
// Therefore, the answer is 2.  

// Example 3:
// Input: words = ["abab","ab"]
// Output: 0
// Explanation: In this example, the only valid index pair is i = 0 and j = 1, and isPrefixAndSuffix("abab", "ab") is false.
// Therefore, the answer is 0.

// Constraints:
//     1 <= words.length <= 10^5
//     1 <= words[i].length <= 10^5
//     words[i] consists only of lowercase English letters.
//     The sum of the lengths of all words[i] does not exceed 5 * 10^5.

import "fmt"
//import "strings"

// // Time Limit Exceeded 591 / 596
// func countPrefixSuffixPairs(words []string) int64 {
//     res := 0
//     for i, s := range words {
//         for _, t := range words[i+1:] {
//             if strings.HasPrefix(t, s) && strings.HasSuffix(t, s) {
//                 res++
//             }
//         }
//     }
//     return int64(res)
// }


// // Time Limit Exceeded 591 / 596
// func countPrefixSuffixPairs1(words []string) int64 {
//     res := 0
//     for i, w := range words {
//         for j := i + 1; j < len(words); j++ {
//             if len(w) > len(words[j]) { continue }
//             if w == words[j][:len(w)] && w == words[j][len(words[j]) - len(w):] {
//                 res++
//             }
//         }
//     }
//     return int64(res)
// }

type Node struct {
    children map[int]*Node
    count      int
}

func countPrefixSuffixPairs(words []string) int64 {
    res := 0
    trie := &Node{ children: make(map[int]*Node) }
    for _, s := range words {
        node := trie
        m := len(s)
        for i := 0; i < m; i++ {
            p := int(s[i]) * 32 + int(s[m-i-1])
            if _, ok := node.children[p]; !ok {
                node.children[p] = &Node{children: make(map[int]*Node)}
            }
            node = node.children[p]
            res += node.count
        }
        node.count++
    }
    return int64(res)
}

func countPrefixSuffixPairs1(words []string) int64 {
    res := 0
    mp := make(map[string]int)
    max := func (x, y int) int { if x > y { return x; }; return y; }
    for _, s := range words {
        n := len(s)
        arr:= make([]int, n)
        for i, l, r := 0, 0, 0; i < n; i++ {
            arr[i] = max(min(arr[i-l], r - i + 1), 0)
            for i + arr[i] < n && s[arr[i]] == s[i + arr[i]] {
                l, r = i, i + arr[i]
                arr[i]++
            }
            if arr[i] == n - i {
                res += mp[s[:arr[i]]]
            }
        }
        mp[s]++
    }
    return int64(res)
}

func main() {
    // Example 1:
    // Input: words = ["a","aba","ababa","aa"]
    // Output: 4
    // Explanation: In this example, the counted index pairs are:
    // i = 0 and j = 1 because isPrefixAndSuffix("a", "aba") is true.
    // i = 0 and j = 2 because isPrefixAndSuffix("a", "ababa") is true.
    // i = 0 and j = 3 because isPrefixAndSuffix("a", "aa") is true.
    // i = 1 and j = 2 because isPrefixAndSuffix("aba", "ababa") is true.
    // Therefore, the answer is 4.
    fmt.Println(countPrefixSuffixPairs([]string{"a","aba","ababa","aa"})) // 4
    // Example 2:
    // Input: words = ["pa","papa","ma","mama"]
    // Output: 2
    // Explanation: In this example, the counted index pairs are:
    // i = 0 and j = 1 because isPrefixAndSuffix("pa", "papa") is true.
    // i = 2 and j = 3 because isPrefixAndSuffix("ma", "mama") is true.
    // Therefore, the answer is 2.
    fmt.Println(countPrefixSuffixPairs([]string{"pa","papa","ma","mama"})) // 2
    // Example 3:
    // Input: words = ["abab","ab"]
    // Output: 0
    // Explanation: In this example, the only valid index pair is i = 0 and j = 1, and isPrefixAndSuffix("abab", "ab") is false.
    // Therefore, the answer is 0.
    fmt.Println(countPrefixSuffixPairs([]string{"abab","ab"})) // 0

    fmt.Println(countPrefixSuffixPairs1([]string{"a","aba","ababa","aa"})) // 4
    fmt.Println(countPrefixSuffixPairs1([]string{"pa","papa","ma","mama"})) // 2
    fmt.Println(countPrefixSuffixPairs1([]string{"abab","ab"})) // 0
}