Loops optimized for knapsack on Tree - LQT

//Make CSP great again
//You're as beautiful as the day I lost you
//New year, best wishes
#include <bits/stdc++.h>
#define TASK "TESTCODE"
#define Log2(x) 31 - __builtin_clz(x)
using namespace std;
const int N = 3000, mod = 998244353;
int n, a[N + 1], dp[N + 1][2 * N + 1], s[2 * N + 1], sz[N + 1], cnt[N + 1], res = 0;
vector<int> Adj[N + 1];
void read()
{
    cin >> n;
    for (int i = 1; i <= n; ++ i)
    {
        cin >> a[i];
        ++cnt[a[i]];
    }
    for (int i = 1; i < n; ++ i)
    {
        int u, v;
        cin >> u >> v;
        Adj[u].push_back(v);
        Adj[v].push_back(u);
    }
}
int g(int u, int c)
{
    return a[u] == c ? 1 : -1;
}
void DFS(int u, int p, int c)
{
    int t = cnt[c];
    sz[u] = 1;
    for (int i = -t; i <= t; ++ i)
    {
        dp[u][i + n] = 0;
    }
    dp[u][g(u, c) + n] = 1;
    for (int v : Adj[u])
    {
        if (v == p)
        {
            continue;
        }
        DFS(v, u, c);
        fill(s + n - min(t, sz[u] + sz[v]), s + n + min(t, sz[u] + sz[v]) + 1, 0);
        //memset(s, 0, sizeof(s));
        for (int i = -sz[u]; i <= sz[u]; ++ i)
        {
            if (dp[u][i + n] != 0)
            {
                s[i + n] += dp[u][i + n];
                s[i + n] %= mod;
                for (int j = -sz[v]; j <= sz[v]; ++ j)
                {
                    if (i + j < -t || i + j > t)
                    {
                        continue;
                    }
                    if (dp[v][j + n] != 0)
                    {
                        s[i + j + n] = (s[i + j + n] + (1LL * dp[u][i + n] * dp[v][j + n]) % mod) % mod;
                    }
                }
            }
        }
        for (int j = -sz[v]; j <= sz[v]; ++ j)
        {
            dp[v][j + n] = 0;
        }
        sz[u] += sz[v];
        sz[u] = min(sz[u], t);
        for (int i = -sz[u]; i <= sz[u]; ++ i)
        {
            dp[u][i + n] = s[i + n];
        }
    }
    for (int i = 1; i <= min(sz[u], t); ++ i)
    {
        res += dp[u][i + n];
        res %= mod;
    }
}
void solve()
{
    for (int i = 1; i <= n; ++ i)
    {
        DFS(1, 0, i);
    }
    cout << res;
}
int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    if (fopen(TASK".INP", "r"))
    {
        freopen(TASK".INP", "r", stdin);
        //freopen(TASK".out", "w", stdout);
    }
    int t = 1;
    bool typetest = false;
    if (typetest)
    {
        cin >> t;
    }
    for (int __ = 1; __ <= t; ++ __)
    {
        //cout << "Case " << __ << ": ";
        read();
        solve();
    }
}