题解 | #打家劫舍(三)#

package main

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

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 << 3 + ans <<1 + int(c-'0')
        read = true
    }
    
    if read {
        *a = ans *sign
    }
    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)
    cost, edge, dp := make([]int, n+1), make([][]int, n+1), make([][2]int, n+1)
    for i := 1;  i<=n; i++ {
        ReadInt(&cost[i])
    }
    for i := 0;  i<n; i++ {
        var a int
        ReadInt(&a)
        edge[a] = append(edge[a], i+1)
    }

    var dfs func(int, int) 
    dfs = func(root, fa int) {
        dp[root][0] += cost[root]
        for _, item := range edge[root] {
            if item == fa {
                continue
            }
            dfs(item, root)
            dp[root][0] += dp[item][1]
            dp[root][1] += max(dp[item][0], dp[item][1])

        }
    }
    dfs(0, -1)
    ans := max(dp[0][1], dp[0][0])
    fmt.Fprintln(out, ans)
}