← Home
For problem statement at 0-999/800-899/890-899/894/problemE.txt this is a correct solution, but verifier at 0-999/800-899/890-899/894/verifierE.go ends with All tests passed can you fix the verifier? package main

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

type FastScanner struct {
	buf []byte
	idx int
	n   int
}

func NewFastScanner() *FastScanner {
	return &FastScanner{buf: make([]byte, 1<<20)}
}

func (fs *FastScanner) readByte() byte {
	if fs.idx >= fs.n {
		n, _ := os.Stdin.Read(fs.buf)
		if n == 0 {
			return 0
		}
		fs.idx = 0
		fs.n = n
	}
	b := fs.buf[fs.idx]
	fs.idx++
	return b
}

func (fs *FastScanner) NextInt() int {
	sign, val := 1, 0
	c := fs.readByte()
	for c != '-' && (c < '0' || c > '9') {
		c = fs.readByte()
	}
	if c == '-' {
		sign = -1
		c = fs.readByte()
	}
	for c >= '0' && c <= '9' {
		val = val*10 + int(c-'0')
		c = fs.readByte()
	}
	return sign * val
}

func gain(w int64) int64 {
	l, r := int64(0), int64(200000)
	for l < r {
		m := (l + r + 1) >> 1
		if m*(m+1)/2 <= w {
			l = m
		} else {
			r = m - 1
		}
	}
	k := l
	return (k+1)*w - k*(k+1)*(k+2)/6
}

func main() {
	in := NewFastScanner()
	n := in.NextInt()
	m := in.NextInt()

	headF := make([]int, n+1)
	headR := make([]int, n+1)
	for i := 1; i <= n; i++ {
		headF[i] = -1
		headR[i] = -1
	}

	to := make([]int, m)
	revTo := make([]int, m)
	nextF := make([]int, m)
	nextR := make([]int, m)
	wt := make([]int64, m)

	for i := 0; i < m; i++ {
		x := in.NextInt()
		y := in.NextInt()
		w := int64(in.NextInt())
		to[i] = y
		revTo[i] = x
		wt[i] = w
		nextF[i] = headF[x]
		headF[x] = i
		nextR[i] = headR[y]
		headR[y] = i
	}

	s := in.NextInt()

	vis := make([]bool, n+1)
	order := make([]int, 0)
	nodeSt := make([]int, 0)
	edgeSt := make([]int, 0)

	vis[s] = true
	nodeSt = append(nodeSt, s)
	edgeSt = append(edgeSt, headF[s])

	for len(nodeSt) > 0 {
		i := len(nodeSt) - 1
		e := edgeSt[i]
		if e == -1 {
			u := nodeSt[i]
			nodeSt = nodeSt[:i]
			edgeSt = edgeSt[:i]
			order = append(order, u)
			continue
		}
		edgeSt[i] = nextF[e]
		v := to[e]
		if !vis[v] {
			vis[v] = true
			nodeSt = append(nodeSt, v)
			edgeSt = append(edgeSt, headF[v])
		}
	}

	comp := make([]int, n+1)
	nodeSt = nodeSt[:0]
	compCnt := 0

	for i := len(order) - 1; i >= 0; i-- {
		u := order[i]
		if comp[u] != 0 {
			continue
		}
		compCnt++
		comp[u] = compCnt
		nodeSt = append(nodeSt, u)
		for len(nodeSt) > 0 {
			x := nodeSt[len(nodeSt)-1]
			nodeSt = nodeSt[:len(nodeSt)-1]
			for e := headR[x]; e != -1; e = nextR[e] {
				v := revTo[e]
				if vis[v] && comp[v] == 0 {
					comp[v] = compCnt
					nodeSt = append(nodeSt, v)
				}
			}
		}
	}

	compHead := make([]int, compCnt+1)
	nextNode := make([]int, n+1)
	compWeight := make([]int64, compCnt+1)

	for _, u := range order {
		c := comp[u]
		nextNode[u] = compHead[c]
		compHead[c] = u
		for e := headF[u]; e != -1; e = nextF[e] {
			if comp[to[e]] == c {
				compWeight[c] += gain(wt[e])
			}
		}
	}

	dp := make([]int64, compCnt+1)
	for i := 1; i <= compCnt; i++ {
		dp[i] = -1
	}

	src := comp[s]
	dp[src] = compWeight[src]
	ans := dp[src]

	for c := 1; c <= compCnt; c++ {
		if dp[c] < 0 {
			continue
		}
		if dp[c] > ans {
			ans = dp[c]
		}
		for u := compHead[c]; u != 0; u = nextNode[u] {
			for e := headF[u]; e != -1; e = nextF[e] {
				vComp := comp[to[e]]
				if vComp != c {
					val := dp[c] + wt[e] + compWeight[vComp]
					if val > dp[vComp] {
						dp[vComp] = val
					}
				}
			}
		}
	}

	out := bufio.NewWriterSize(os.Stdout, 1<<20)
	fmt.Fprintln(out, ans)
	out.Flush()
}