网易笔试题复盘

T4

题目描述：

子树的权值为子树中所有节点权值乘积的因子数，求所有子树的权值和，取模10^9+7
输入：
第一行为：n
第二行为：每个节点的权值
后n-1行为两点的连接
3
1 2 3
1 2
1 3

输出：8

思路：

做这种题，我的思路都是先找叶子节点，然后从叶子节点往根节点上去遍历。原因是父节点需要用到子节点的信息，所以应该先计算子节点，再去计算父节点。可以用一个小技巧去找到所有的叶子节点，就是没有指向其他任何节点的节点，也就是在边 edges 数组中没有在edges[i][0]种出现过的节点。

找到所有的叶子节点之后，依次遍历每一条从叶子节点到根节点的路径。其中需要使用used数组去重，避免某个节点的权值被乘了两次；

最后使用一个 score 数组去记录遍历到每一个节点时的值；

以下代码是在考试结束之后才调出来的，未验证正确性，有大佬看到错误还请告知！

#include <iostream>
#include <vector>
#include <cmath>
#include <unordered_set>
#include <unordered_map>
using namespace std;

int yinzi_num(int n) {
    int res = 0;
    for (int i = 1; i <= sqrt(n); ++i) {
        if (n % i == 0) {
            res += (i == n / i) ? 1 : 2;
            res %= 1000000007;
        }
    }
    return res;
}

int main() {
    int n;
    cin >> n;
    /* 每个节点权值 */
    vector<int> nums(n);
    for (int i = 0; i < n; ++i) {
        cin >> nums[i];
    }
    
    /* n - 1 条边 */
    vector<vector<int>> edges(n - 1, vector<int>(2));
    for (int i = 0; i < n - 1; ++i) {
        cin >> edges[i][0] >> edges[i][1];
    }
    
    /* 寻找叶子节点 */
    vector<int> yezi;
    unordered_set<int> hashset;
    for (auto v : edges) {
        hashset.insert(v[0]);
    }
    for (int i = 1; i <= n; ++i) {
        if (hashset.find(i) == hashset.end()) yezi.push_back(i);
    }
    
    /* 记录每个节点的父节点 */
    unordered_map<int, int> hashmap;
    hashmap[1] = 0;
    for (auto v : edges) {
        hashmap[v[1]] = v[0];
    }
    
    vector<int> score(n + 1, 1);
    vector<bool> used(n + 1, false);
    for (int i = 0; i < yezi.size(); ++i) {
        int sum = 1;
        int node = yezi[i];
        while (node != 0) {
            sum *= (used[node] ? 1 : nums[node - 1]);
            score[node] *= sum;
            used[node] = true;
            node = hashmap[node];
        }
    }
    
    int res = 0;
    for (int i = 1; i <= n; ++i) {
        res += yinzi_num(score[i]);
        res %= 1000000007;
    }
    
    cout << res << endl;
    
    return 0;
}