题解 | 喜欢切数组的红

package main

import (
	"fmt"
)

func main() {
    var n int
    fmt.Scan(&n)
    nums := make([]int, n)
    sum := 0
    for i:=0; i< n; i++ {
        fmt.Scan(&nums[i])
        sum += nums[i]
    }
    
    if sum%3 != 0 {
        fmt.Println(0)
        return
    }
    target := sum/3
    preSum := make([]int,n)
    posCount := make([]int, n)
    posTotal := 0
    sumTotal := 0
    for i:=0; i <n; i++ {
        sumTotal += nums[i]
        preSum[i] = sumTotal
        if nums[i] > 0 {
            posTotal++
        }
        posCount[i] = posTotal
    }

    count := 0
    firstCuts :=make(map[int]struct{})

    for j:=1; j < n-1; j++ {
        if preSum[j-1] == target && posCount[j-1] >0 {
            firstCuts[j-1] = struct{}{} 
        }

        if preSum[j] == 2 * target && posCount[j] >0 {
            for i := range firstCuts {
                if posCount[i] < posCount[j] { 
                    count++
                }
            } 
        }
    }

    fmt.Println(count)

}