CF1101 D GCD Counting

犯了一个愚蠢的错误。

vb vis(N);
vi a(N);
vv<int> p(N);
vv<int> g(N);


vector<map<int,int>> f(N);
int ans = 0;

void dfs(int u) {
    vis[u] = true;
    int t = a[u];

    map<int,int> b1, b2;

    FOR(x, p[t]) {
        f[u][x] = 1;
    }

    FOR(v, g[u]) {
        if(!vis[v]) {
            dfs(v);
            FOR(x, p[t]) {
                auto iter = f[v].find(x);
                if(iter != f[v].end()) {
                    int tmp = iter->second;
                    if (b1[x] < tmp) {
                        b2[x] = b1[x];
                        b1[x] = tmp;
                    }
                    else if(b2[x] < tmp) {
                        b2[x] = tmp;
                    }
                }
            }
        }
    }

    FOR(x, p[t]) {
         f[u][x] += b1[x] + b2[x];
         ans = max(ans, f[u][x]);
    }
}

int main() {
    FAST_READ
#ifdef LOCAL
    ifstream in("main.in");
    cin.rdbuf(in.rdbuf());
#endif

    const int maxn = 2e5+5;

    vb isp(maxn, true);

    rng(i, 2, maxn) {
        if(isp[i]) {
            p[i].pb(i);
            stp(j, 2*i, maxn, i) {
                isp[j] = false;
                p[j].pb(i);
            }
        }
    }

    int n; cin >> n;

  rng(i, 1, n+1) {
        cin >> a[i];
    }

    rep(n-1) {
        int u, v;
        scan(u, v);
        if(__gcd(a[u], a[v]) >1) {
            g[u].pb(v);
            g[v].pb(u);
        }
    }

    rng(i, 1, n + 1) {
      if(!vis[i]) dfs(i);
  }

  println(ans)

#ifdef LOCAL
    cout << "Time elapsed: " << 1.0 * clock() / CLOCKS_PER_SEC << " s.\n";
#endif
    return 0;
}

在 dfs 函数的最后

    FOR(x, p[t]) {
         f[u][x] += b1[x] + b2[x];
         ans = max(ans, f[u][x]);
    }

写错了，应该是

FOR (x, p[t]) {
  f[u][x] += b1[x];
  ans = max(ans, f[u][x] + b2[x]);
} 

诶，搞不懂我为何如此傻逼。
我看了 palayutm 的 submission，发现我的维护最大和次大的写法实在是太蠢了。实际上这类树上 DP 问题都不必这样写。dfs 函数比较好的写法是

void dfs(int u) {
    vis[u] = true;
    int t = a[u];

    map<int,int> b1, b2;

    FOR(x, p[t]) {
        f[u][x] = 1, ans = max(ans, 1);
    }

    FOR(v, g[u]) {
        if(!vis[v]) {
            dfs(v);
            FOR(x, p[t]) {
                auto iter = f[v].find(x);
                if(iter != f[v].end()) {
                    int tmp = iter->second;
                    ans = max(ans, f[u][v] + tmp);
                    f[u][v] = max(f[u][v], 1 + tmp);
                }
            }
        }
    }
}

另一个应该改进的写法是求 的质因子。自然是用埃氏筛。质因子放到 std::vector<int> 里。比较紧凑的写法是

// vector<vector<int>> p(N);
rng (i, 2, SZ(p)) if (p[i].empty()) 
    for (int j = i; j < SZ(p); j += i) p[j].pb(i);