fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. using ll = long long;
  5.  
  6. const int MAXN = 5005;
  7. const int MAXK = 6;
  8. const int MAXMASK = 1 << MAXK;
  9. const int MOD = 998244353;
  10.  
  11. int n, k;
  12. vector<int> g[MAXN];
  13.  
  14. int s[MAXK + 2];
  15. int sz[MAXN];
  16.  
  17. // mx[mask] = największy numer operacji występujący w masce
  18. int mx[MAXMASK];
  19.  
  20. // dp[v][first][mask]
  21. //
  22. // first:
  23. // 1..k   -> najwcześniejsza operacja, która jest typu "keep subtree"
  24. // k+1    -> nie ma jeszcze żadnej operacji "keep subtree"
  25. //
  26. int dp[MAXN][MAXK + 2][MAXMASK];
  27.  
  28. void dfs(int v, int par) {
  29.     sz[v] = 1;
  30.  
  31.     // Na początku poddrzewo v zawiera tylko v.
  32.     dp[v][k + 1][0] = 1;
  33.  
  34.     for (int u : g[v]) {
  35.         if (u == par) continue;
  36.  
  37.         dfs(u, v);
  38.         sz[v] += sz[u];
  39.  
  40.         static int tmp[MAXK + 2][MAXMASK];
  41.  
  42.         for (int i = 1; i <= k + 1; ++i)
  43.             for (int mask = 0; mask < (1 << k); ++mask)
  44.                 tmp[i][mask] = 0;
  45.  
  46.         // Scalanie poddrzewa u z v.
  47.         for (int first = 1; first <= k + 1; ++first) {
  48.             for (int mask = 0; mask < (1 << k); ++mask) {
  49.  
  50.                 if (!dp[v][first][mask]) continue;
  51.  
  52.                 int available = (1 << k) - 1 - mask;
  53.  
  54.                 // subset = operacje pochodzące z u
  55.                 for (int subset = available;; subset = (subset - 1) & available) {
  56.  
  57.                     // u nie ma własnego "keep subtree"
  58.                     // albo jego pierwsze keep jest późniejsze niż first.
  59.                     if (mx[subset] <= first) {
  60.                         tmp[first][mask | subset] =
  61.                             (tmp[first][mask | subset]
  62.                              + 1LL * dp[v][first][mask]
  63.                              * dp[u][k + 1][subset]) % MOD;
  64.                     }
  65.  
  66.                     // v nie ma własnego "keep subtree"
  67.                     // albo jego pierwsze keep jest późniejsze/nie wcześniejsze.
  68.                     if (first != k + 1 && mx[mask] <= first) {
  69.                         tmp[first][mask | subset] =
  70.                             (tmp[first][mask | subset]
  71.                              + 1LL * dp[v][k + 1][mask]
  72.                              * dp[u][first][subset]) % MOD;
  73.                     }
  74.  
  75.                     if (subset == 0) break;
  76.                 }
  77.             }
  78.         }
  79.  
  80.         for (int first = 1; first <= k + 1; ++first)
  81.             for (int mask = 0; mask < (1 << k); ++mask)
  82.                 dp[v][first][mask] = tmp[first][mask];
  83.     }
  84.  
  85.     // Krawędź (v, parent).
  86.     // Dla korzenia jej nie ma.
  87.     if (v == 1) return;
  88.  
  89.     static int tmp[MAXK + 2][MAXMASK];
  90.  
  91.     for (int first = 1; first <= k + 1; ++first)
  92.         for (int mask = 0; mask < (1 << k); ++mask)
  93.             tmp[first][mask] = dp[v][first][mask];
  94.  
  95.     for (int first = 1; first <= k + 1; ++first) {
  96.         for (int mask = 0; mask < (1 << k); ++mask) {
  97.  
  98.             if (!dp[v][first][mask]) continue;
  99.  
  100.             // Wybieramy, że krawędź (v,parent)
  101.             // odpowiada operacji j.
  102.             for (int j = 1; j < first; ++j) {
  103.  
  104.                 if (mask & (1 << (j - 1)))
  105.                     continue;
  106.  
  107.                 // Rozmiar aktualnego komponentu w poddrzewie v.
  108.                 // Wcześniejsze operacje < j, które są w tej masce,
  109.                 // musiały usunąć fragmenty.
  110.                 int currentSize = sz[v];
  111.  
  112.                 for (int t = 1; t < j; ++t) {
  113.                     if (mask & (1 << (t - 1))) {
  114.                         currentSize -= s[t - 1] - s[t];
  115.                     }
  116.                 }
  117.  
  118.                 // 1) Operacja j = "keep subtree"
  119.                 //
  120.                 // Po odcięciu krawędzi zostawiamy stronę zawierającą v.
  121.                 // Jej rozmiar musi być s[j].
  122.                 if (currentSize == s[j]) {
  123.                     tmp[j][mask | (1 << (j - 1))] =
  124.                         (tmp[j][mask | (1 << (j - 1))]
  125.                          + dp[v][first][mask]) % MOD;
  126.                 }
  127.  
  128.                 // 2) Operacja j = "delete subtree"
  129.                 //
  130.                 // Usuwamy stronę zawierającą v.
  131.                 // Musi mieć rozmiar s[j-1] - s[j].
  132.                 //
  133.                 // Wszystkie operacje w tej części muszą być wcześniejsze
  134.                 // i żadna z nich nie może być typu "keep subtree".
  135.                 if (currentSize == s[j - 1] - s[j]
  136.                     && mx[mask] < j) {
  137.  
  138.                     tmp[first][mask | (1 << (j - 1))] =
  139.                         (tmp[first][mask | (1 << (j - 1))]
  140.                          + dp[v][first][mask]) % MOD;
  141.                 }
  142.             }
  143.         }
  144.     }
  145.  
  146.     for (int first = 1; first <= k + 1; ++first)
  147.         for (int mask = 0; mask < (1 << k); ++mask)
  148.             dp[v][first][mask] = tmp[first][mask];
  149. }
  150.  
  151. int main() {
  152.     ios::sync_with_stdio(false);
  153.     cin.tie(nullptr);
  154.  
  155.     cin >> n;
  156.  
  157.     for (int i = 0; i < n - 1; ++i) {
  158.         int a, b;
  159.         cin >> a >> b;
  160.         g[a].push_back(b);
  161.         g[b].push_back(a);
  162.     }
  163.  
  164.     cin >> k;
  165.  
  166.     s[0] = n;
  167.     for (int i = 1; i <= k; ++i) {
  168.         cin >> s[i];
  169.     }
  170.  
  171.     // Największy numer operacji w masce.
  172.     for (int mask = 1; mask < (1 << k); ++mask) {
  173.         for (int i = 1; i <= k; ++i) {
  174.             if (mask & (1 << (i - 1))) {
  175.                 mx[mask] = max(mx[mask], i);
  176.             }
  177.         }
  178.     }
  179.  
  180.     dfs(1, 0);
  181.  
  182.     int all = (1 << k) - 1;
  183.  
  184.     ll ans = 0;
  185.  
  186.     for (int first = 1; first <= k + 1; ++first) {
  187.         ans += dp[1][first][all];
  188.         ans %= MOD;
  189.     }
  190.  
  191.     cout << ans << '\n';
  192.  
  193.     return 0;
  194. }
Success #stdin #stdout 0.01s 5292KB
comments (?)
stdin
copy 
3
1 2
2 3
2
2 1
stdout
copy 
1
https://ideone.com/3BNLx3
language:
C++ (gcc 8.3)
created:
8 hours ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!