package main

func candy(ratings []int) (sum int) {
	n := len(ratings)
	if n <= 1 {
		return n
	}
	candies := make([]int, n)
	for i := range candies {
		candies[i] = 1
	}
	for i := 1; i < n; i++ { // Forward increasing
		if ratings[i-1] < ratings[i] {
			candies[i] = candies[i-1] + 1
		}
	}
	for i := n - 1; 0 < i; i-- { // Backward increasing
		if ratings[i] < ratings[i-1] && candies[i-1] < candies[i]+1 {
			candies[i-1] = candies[i] + 1
		}
		sum += candies[i]
	}
	return sum + candies[0]
}

// func main() {
// 	println(candy([]int{1, 0, 2}), 5)
// 	println(candy([]int{1, 2, 2}), 4)
// }