题解 | #二叉树中的最大路径和#

package main

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

const MIN = -1000000000

var in = bufio.NewReader(os.Stdin)
var out = bufio.NewWriter(os.Stdout)

func readInt(a *int) error {
	sign, ans, read := 1, 0, false
	c, err := in.ReadByte()

	for ; err == nil && (c < '0' || c > '9'); c, err = in.ReadByte() {
		if c == '-' {
			sign = -1
		}
	}
	for ; err == nil && c >= '0' && c <= '9'; c, err = in.ReadByte() {
		ans = ans *10 + int(c-'0')
        read = true
	}
    if read {
        *a = ans * sign
        return nil
    }
    return err
}

func max(a, b int) int {
    if a< b {
        return b
    }
    return a
}
func main() {
    defer out.Flush()
	var n int
    readInt(&n)
    var cost = make([]int, n+1)
    var edge = make([][]int, n+1)
    for i := 1; i <= n; i++ {
        readInt(&cost[i])
    }
    root := 0
    for i := 1; i <= n; i++ {
        var cur int
        readInt(&cur)
        if cur == 0 {
            root = i
        } else {
            edge[cur] = append(edge[cur], i)
        }
    }
    ans := MIN
    var dfs func(cur, fa int) int 
    dfs = func(cur, fa int) int{
        ans = max(ans, cost[cur])
        var left, right int= MIN, MIN
        var i int = 0
        for _, item := range edge[cur] {
            if item == fa {
                continue
            }
            if i == 0 {
                left = dfs(item, cur)
            } else {
                right = dfs(item, cur)
            }
            i++ 
        }
        
        ans = max(ans, cost[cur]+max(left+right, max(left, right)))
        return max(cost[cur], cost[cur]+max(left, right))
    }
    dfs(root, -1)
    fmt.Fprintln(out, ans)
}