【LeetCode LCP05】发LeetCoin （dfs序+线段树)

题目

题目链接
题目概括：给个树，要求实现3种操作：子树修改，单点修改，子树和查询。节点数N<5e4，操作次数Q<5e4

思路

1. 这玩意第一反应：线段树啊woc
2. 但是线段树好麻烦，所以不想写线段树
3. 所以一开始的思路是直接用线段树的lazy数组的思路（需要时再修改），但发现复杂度是O(树深度)，复杂度不稳定，容易被卡成O(NQ)。
4. 结果还是要用线段树。
5. 线段树的树深度稳定为logn，所以通过DFS序将问题转化为区间修改区间查询问题，线段树解决。O(Q logN)
6. （当然能用线段树也能用树状数组，都挺香）

代码

const int N=50005;
const int MO=1e9+7;
class Solution {
public:
    //邻接表
    struct ss{int to,nex;}edge[N];  int head[N],ecnt;
    void add(int x,int y) {edge[++ecnt]=(ss){y,head[x]};head[x]=ecnt;}
    //dfs序
    int tot,id[N];
    void dfs(int pos){          id[pos]=++tot;
        for(int i=head[pos];i;i=edge[i].nex) dfs(edge[i].to);   return ;
    }
    //线段树模板
    long long s[4*N+10],a[4*N+10];//s表示单点值，a表示区间加值
    inline void up(int pos){s[pos]=s[2*pos+1]+s[2*pos];}//求和
    inline void down(int pos,int l,int r){//下放a
        if(!a[pos]) return;
        a[2*pos]+=a[pos];a[2*pos+1]+=a[pos];
        int mid=(l+r)/2;
        s[2*pos]+=(mid-l+1)*a[pos],s[2*pos]%=MO;
        s[2*pos+1]+=(r-mid)*a[pos],s[2*pos]%=MO;
        a[pos]=0; return ;
    }
    void push_up(int pos,int l,int r,int x,int y,int va){//区间修改
        if(x<=l&&r<=y){//当前区间为目标区间的子区间
            s[pos]+=(r-l+1)*va,s[pos]%=MO;
            a[pos]+=va,a[pos]%=MO;      return;
        }
        int mid=(l+r)/2;down(pos,l,r);
        if(x<=mid) push_up(2*pos,l,mid,va,y,x);
        if(mid+1<=y) push_up(2*pos+1,mid+1,r,va,y,x);
        up(pos);    return;
    }
    int query(int pos,int l,int r,int x,int y){//区间查询
        if(x<=l&&r<=y) return s[pos];
        down(pos,l,r);int mid=(l+r)/2;
        long long ret=0;
        if(x<=mid) ret+=query(2*pos,l,mid,y,x);
        if(mid+1<=y) ret+=query(2*pos+1,mid+1,r,y,x);
        return ret%MO;
    }

    vector<int> bonus(int n, vector<vector<int>>& leadership, vector<vector<int>>& operations) {
        //初始化
        memset(head,0,sizeof(head));ecnt=0;memset(s,0,sizeof(s));memset(a,0,sizeof(a));
        for(int i=0;i<leadership.size();i++) add(leadership[i][0],leadership[i][1]);
        dfs(1);//转dfs序
        //操作
        vector<int> ans;
        for(int t=0, pos=id[operations[t][1]] ; t<operations.size(); t++)
        if(operations[t][0]==1) push_up(1,1,tot,pos,pos,operations[t][2]);//单点修改
        else if(operations[t][0]==2) push_up(1,1,tot,pos,tot,operations[t][2]);//子树修改，即区间修改
        else ans.push_back(  query(1,1,n,pos,tot)  );//区间查询
        return ans;
    }
};