CF894E Ralph and Mushrooms

#include<bits/stdc++.h>
#define rg register
#define ll long long
using namespace std;
const int N=1e6+5;
const ll INF=1ll<<60;
int tot,head[N],nxt[N],ver[N];
int n,m,s,x[N],y[N];
int num,dfn[N],low[N];
int top,st[N],col,co[N];
int dgr[N];
ll val[N],f[N],w[N],c[N];
inline int read() {
    int X=0,w=0; char ch=0;
    while(!isdigit(ch)) {w|=ch=='-';ch=getchar();}
    while(isdigit(ch)) X=(X<<3)+(X<<1)+(ch^48),ch=getchar();
    return w?-X:X;
}
inline void addedge(int u,int v,ll d) {
    nxt[++tot]=head[u],head[u]=tot,ver[tot]=v,w[tot]=d;
}

inline void tarjan(int u) {
    dfn[u]=low[u]=++num,st[++top]=u;
    for (rg  int i=head[u];i;i=nxt[i]) {
        int v=ver[i];
        if (!dfn[v])
            tarjan(v),low[u]=min(low[u],low[v]);
        else
            if (!co[v])low[u]=min(low[u],dfn[v]);
    }
    if (low[u]==dfn[u]) {
        ++col;
        do co[st[top]]=col;
        while (st[top--]!=u);
    }
}

inline ll calc(ll c) {
    ll t=sqrt(c*2+0.25)-0.5;
    return (t+1)*c-t*(t+1)*(t+2)/6;
}

inline void rebuild() {
    for (rg int i=1;i<=m;++i)
        if (co[x[i]]==co[y[i]])
            val[co[x[i]]]+=calc(c[i]);
    memset(head,tot=0,sizeof head);    
    for (rg int i=1;i<=m;++i)
        if (co[x[i]]!=co[y[i]])
            addedge(co[x[i]],co[y[i]],c[i]+val[co[y[i]]]),++dgr[co[y[i]]];
}

inline ll toposort() {
    ll ans=0;
    queue<int >Q;
    for (rg int i=1;i<=col;++i) {
        if (dgr[i]==0)Q.push(i);f[i]=-INF;
    }
    f[co[s]]=val[co[s]];
    while (!Q.empty()) {
        int u=Q.front();Q.pop();
        for (rg int v,i=head[u];i;i=nxt[i]) {
            if (!--dgr[v=ver[i]])Q.push(v);
            f[v]=max(f[v],f[u]+w[i]);
        }
    }
    for (rg int i=1;i<=col;++i) ans=max(ans,f[i]);
    return ans;
}

int main() {
    n=read(),m=read();
    for (rg int i=1;i<=m;++i)
        x[i]=read(),y[i]=read(),c[i]=read(),addedge(x[i],y[i],c[i]);
    s=read();
    for (rg int i=1;i<=n;++i)
        if (!dfn[i]) tarjan(i);
    rebuild();
    printf("%lld
",toposort());
    return 0;
}