LeetcodeLCP 05-发 LeetCoin

整个员工关系为一棵树，我们求出它的DFS序，就可以把树上的修改操作转化为区间操作，然后用线段树来维护区间信息即可。

class Solution {
public:
    static const int mod=1e9+7,N=50010;
    struct Segment{
        long long l,r,v,sum;
    }tr[N*4];
    void build(int u,long long l,long long r){
        tr[u]={l,r,0,0};
        if(l==r) return;
        long long m=(l+r)>>1;
        build(u<<1,l,m);
        build(u<<1|1,m+1,r);
    }
    void pushup(int u){
        tr[u].v=(tr[u<<1].v+tr[u<<1|1].v)%mod;
    }
    void pushdown(int u){
        if(!tr[u].sum) return;
        tr[u<<1].v=(tr[u<<1].v+tr[u].sum*(tr[u<<1].r-tr[u<<1].l+1)%mod)%mod;
        tr[u<<1|1].v=(tr[u<<1|1].v+tr[u].sum*(tr[u<<1|1].r-tr[u<<1|1].l+1)%mod)%mod;
        tr[u<<1].sum=(tr[u<<1].sum+tr[u].sum)%mod;
        tr[u<<1|1].sum=(tr[u<<1|1].sum+tr[u].sum)%mod;
        tr[u].sum=0;
    }
    void modify(int u,long long l,long long r,long long k){
        if(tr[u].l>=l&&tr[u].r<=r){
            tr[u].v=(tr[u].v+k*(tr[u].r-tr[u].l+1)%mod)%mod;
            tr[u].sum=(tr[u].sum+k)%mod;
            return;
        }
        pushdown(u);
        long long m=(tr[u].l+tr[u].r)>>1;
        if(l<=m) modify(u<<1,l,r,k);
        if(r>m) modify(u<<1|1,l,r,k);
        pushup(u);
    }
    long long query(int u,long long l,long long r){
        if(tr[u].l>=l&&tr[u].r<=r){
            return tr[u].v; 
        }
        pushdown(u);
        long long res=0;
        long long m=(tr[u].l+tr[u].r)>>1;
        if(l<=m) res=(res+query(u<<1,l,r))%mod;
        if(r>m) res=(res+query(u<<1|1,l,r))%mod;
        return res;
    }
    vector<int> bonus(int n,vector<vector<int>>& leadership, vector<vector<int>>& operations){
        //邻接表建图
        vector<vector<int>> g(n);
        for(auto &p:leadership){
            int x=p[0]-1,y=p[1]-1;
            g[x].push_back(y);
        }
        //求DFS序
        vector<pair<int,int>> hash(n,{0,0});
        int time=0;
        function<void(int)> dfs=[&](int x){
            hash[x].first=++time;
            for(auto y:g[x]) dfs(y);
            hash[x].second=time;
        };
        dfs(0);
        //建立线段树
        build(1,1,n);
        //执行操作
        vector<int> ans;
        for(auto &q:operations){
            int x=q[1]-1;
            if(q[0]==1){
                auto [l,r]=hash[x];
                modify(1,l,l,q[2]);
            }else if(q[0]==2){
                auto [l,r]=hash[x];
                modify(1,l,r,q[2]);
            }else{
                auto [l,r]=hash[x];
                ans.push_back(query(1,l,r));
            }
        }
        return ans;
    }
};

时间复杂度：O（nlogn），n为员工数量。

空间复杂度：O（n）。