整个员工关系为一棵树,我们求出它的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)。