For problem statement at 0-999/800-899/860-869/868/problemE.txt this is a correct solution, but verifier at 0-999/800-899/860-869/868/verifierE.go ends with case 6 failed: expected 9 got 7
input:
4
1 2 4
2 3 3
2 4 1
3
1
2
exit status 1 can you fix the verifier? package main
import (
"bufio"
"fmt"
"os"
"sort"
)
func main() {
in := bufio.NewReader(os.Stdin)
var n int
if _, err := fmt.Fscan(in, &n); err != nil {
return
}
type Edge struct {
to, w int
}
adj := make([][]Edge, n+1)
for i := 0; i < n-1; i++ {
var u, v, w int
fmt.Fscan(in, &u, &v, &w)
adj[u] = append(adj[u], Edge{v, w})
adj[v] = append(adj[v], Edge{u, w})
}
var s int
fmt.Fscan(in, &s)
var m int
fmt.Fscan(in, &m)
criminals := make([]int, m)
for i := 0; i < m; i++ {
fmt.Fscan(in, &criminals[i])
}
active := 0
startCounts := make([]int, n+1)
for _, x := range criminals {
if x != s {
active++
startCounts[x]++
}
}
if active == 0 {
fmt.Println(0)
return
}
parent := make([]int, n+1)
weight := make([]int, n+1)
var order []int
var buildTree func(int, int)
buildTree = func(u, p int) {
order = append(order, u)
for _, e := range adj[u] {
if e.to != p {
parent[e.to] = u
weight[e.to] = e.w
buildTree(e.to, u)
}
}
}
buildTree(s, 0)
subtreeCounts := make([]int, n+1)
for i := n - 1; i >= 0; i-- {
u := order[i]
subtreeCounts[u] = startCounts[u]
for _, e := range adj[u] {
if e.to != parent[u] {
subtreeCounts[u] += subtreeCounts[e.to]
}
}
}
dp0 := make([][]int, n+1)
dp1 := make([][]int, n+1)
for i := 0; i <= n; i++ {
dp0[i] = make([]int, active+1)
dp1[i] = make([]int, active+1)
}
for i := n - 1; i >= 0; i-- {
u := order[i]
if u == s {
continue
}
var children []int
for _, e := range adj[u] {
if e.to != parent[u] {
children = append(children, e.to)
}
}
if len(children) == 0 {
for k := 1; k <= active; k++ {
dp0[u][k] = 0
dp1[u][k] = 0
}
continue
}
cost0 := func(v, c int) int {
if c == 0 {
return 0
}
return 2*weight[v] + dp0[v][c]
}
gain := func(v, c int) int {
if c == 0 {
return 0
}
return cost0(v, c) - weight[v] - dp1[v][c]
}
f := make([]int, active+1)
for k := 1; k <= active; k++ {
f[k] = -1e9
}
for _, v := range children {
nextF := make([]int, active+1)
for k := 0; k <= active; k++ {
nextF[k] = -1e9
}
for k := 0; k <= active; k++ {
if f[k] < 0 {
continue
}
for c := 0; c <= active-k; c++ {
val := f[k] + cost0(v, c)
if val > nextF[k+c] {
nextF[k+c] = val
}
}
}
f = nextF
}
for k := 1; k <= active; k++ {
dp0[u][k] = f[k]
}
var gains []int
for _, v := range children {
for c := 1; c <= active; c++ {
gains = append(gains, gain(v, c))
}
}
sort.Ints(gains)
uniqueGains := []int{}
for j, g := range gains {
if j == 0 || g != gains[j-1] {
uniqueGains = append(uniqueGains, g)
}
}
for k := 1; k <= active; k++ {
dp1[u][k] = -1e9
}
for _, G := range uniqueGains {
f := make([]int, active+1)
for k := 1; k <= active; k++ {
f[k] = -1e9
}
for _, v := range children {
nextF := make([]int, active+1)
for k := 0; k <= active; k++ {
nextF[k] = -1e9
}
for k := 0; k <= active; k++ {
if f[k] < 0 {
continue
}
for c := 0; c <= active-k; c++ {
if c > 0 && gain(v, c) > G {
continue
}
val := f[k] + cost0(v, c)
if val > nextF[k+c] {
nextF[k+c] = val
}
}
}
f = nextF
}
for k := 1; k <= active; k++ {
if f[k] >= 0 {
cand := f[k] - G
if cand > dp1[u][k] {
dp1[u][k] = cand
}
}
}
}
}
var sChildren []int
for _, e := range adj[s] {
sChildren = append(sChildren, e.to)
}
cost0 := func(v, c int) int {
if c == 0 {
return 0
}
return 2*weight[v] + dp0[v][c]
}
gain := func(v, c int) int {
if c == 0 {
return 0
}
return cost0(v, c) - weight[v] - dp1[v][c]
}
getF := func(skipIdx, k int) int {
if k == 0 {
return 0
}
var gains []int
for j, v := range sChildren {
if j == skipIdx {
continue
}
for c := 1; c <= k; c++ {
gains = append(gains, gain(v, c))
}
}
sort.Ints(gains)
uniqueGains := []int{}
for j, g := range gains {
if j == 0 || g != gains[j-1] {
uniqueGains = append(uniqueGains, g)
}
}
ans := int(-1e9)
for _, G := range uniqueGains {
f := make([]int, k+1)
for i := 1; i <= k; i++ {
f[i] = -1e9
}
for j, v := range sChildren {
if j == skipIdx {
continue
}
nextF := make([]int, k+1)
for i := 0; i <= k; i++ {
nextF[i] = -1e9
}
for i := 0; i <= k; i++ {
if f[i] < 0 {
continue
}
for c := 0; c <= k-i; c++ {
if c > 0 && gain(v, c) > G {
continue
}
val := f[i] + cost0(v, c)
if val > nextF[i+c] {
nextF[i+c] = val
}
}
}
f = nextF
}
if f[k] >= 0 {
cand := f[k] - G
if cand > ans {
ans = cand
}
}
}
return ans
}
if len(sChildren) == 0 {
fmt.Println(0)
return
}
if len(sChildren) == 1 {
v := sChildren[0]
fmt.Println(weight[v] + dp1[v][active])
return
}
ans := int(1e9)
for i, v := range sChildren {
c := subtreeCounts[v]
var cand int
if c == active {
cand = weight[v] + dp1[v][active]
} else {
cand = cost0(v, c) + getF(i, active-c)
}
if cand < ans {
ans = cand
}
}
fmt.Println(ans)
}