【C++ | PTA】关于深度优先搜索和逆序对的题应该不会很难吧这件事

题目要求

背景知识
深度优先搜索与 DFS 序
深度优先搜索算法（DFS）是一种用于遍历或搜索树或图的算法。以下伪代码描述了在树 T 上进行深度优先搜索的过程：

procedure DFS(T, u, L)      // T 是被深度优先搜索的树
                            // u 是当前搜索的节点
                            // L 是一个链表，保存了所有节点被第一次访问的顺序
  append u to L             // 将节点 u 添加到链表 L 的末尾
  for v in u.children do    // 枚举节点 u 的所有子节点 v
    DFS(T, v)               // 递归搜索节点 v

令 r 为树 T 的根，调用 DFS(T, r, L) 即可完成对 T 的深度优先搜索，保存在链表 L 中的排列被称为 DFS 序。相信聪明的你已经发现了，如果枚举子节点的顺序不同，最终得到的 DFS 序也会不同。

逆序对
给定一个长度为 n 的整数序列 a1​,a2​,⋯,an​，该序列的逆序对数量是同时满足以下条件的有序数对 (i,j) 的数量：

1≤i<j≤n；
ai​>aj​。

问题求解
给定一棵 n 个节点的树，其中节点 r 为根。求该树所有可能的 DFS 序中逆序对数量之和。
输入格式
第一行输入两个整数 n，r（2≤n≤3×105，1≤r≤n）表示树的大小与根节点。
对于接下来的 (n−1) 行，第 i 行输入两个整数 ui​ 与 vi​（1≤ui​,vi​≤n），表示树上有一条边连接节点 ui​ 与 vi​。

输出格式
输出一行一个整数，表示该树所有可能的 DFS 序中逆序对数量之和。由于答案可能很大，请对 109+7 取模后输出。

样例输入 1
5 3
1 5
2 5
3 5
4 3

样例输出 1
24

样例输入 2
10 5
10 2
2 5
10 7
7 1
7 9
4 2
3 10
10 8
3 6

样例输出 2
516

---

代码

#include<iostream>
#include<vector>
using namespace std;

typedef long long LL;
const int N = 300010, P = 1e9+7;

vector<int> g[N]; // 邻接表存储整张树 
int sz[N],tr[N]; // sz数组记录子树大小，tr数组即为树状数组 
int n,root;
int sum=1,s1,s2; // sum为阶乘，s1和s2用于统计不合法路径的数量 

void add(int x,int y)
{
    for(int i=x;i<N;i+=(i&-i))
        tr[i]+=y;
}

int query(int x)
{
    int res=0;
    for(int i=x;i;i-=(i&-i))
        res+=tr[i];
    return res;
}

void dfs(int u,int fa)
{
    add(u,1); // 将节点u对应位置的计数加1（表示u出现了一次） 
    s1=(s1+query(n)-query(u))%P; // 统计所有深度比u小的节点的DFS序总和 

    sz[u]=1; // 初始化子树大小为1
    int cnt=0; // cnt记录u的子节点个数
    for(auto &j:g[u])
    {
        if(j==fa) continue;
        dfs(j,u); // 递归搜索j的子树 
        sz[u]+=sz[j]; // 更新sz数组
        cnt++;
    } 
    
    for(int i=1;i<=cnt;i++)
        sum=(LL)sum*i%P; // 更新阶乘，即(n-sz[u])! 的值

    s2=(s2+n-query(n)-sz[u]+1)%P; // 统计DFS序比u大的节点个数
    add(u,-1); // 回溯时将u的计数减1（表示u已经不再访问） 
}

int main()
{
    scanf("%d%d",&n,&root);
    for(int i=0;i<n-1;i++)
    {
        int a,b;
        scanf("%d%d",&a,&b);
        g[a].push_back(b);
        g[b].push_back(a); // 无向图，需要双向添加边 
    }

    dfs(root,-1);

    int ans=((LL)s1*sum+(LL)s2*sum%P*(P+1)/4)%P; 
    printf("%d\n",ans);
    return 0; 
}