Salvens

C

Jul 31st, 2023
902
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 2.00 KB | None | 0 0
  1. #include <iostream>
  2. #include <iomanip>
  3. #include <vector>
  4. #include <array>
  5. #include <set>
  6.  
  7. using namespace std;
  8.  
  9. #define int long long
  10.  
  11. const long long INF = 1e18 + 7;
  12. const int MAXN = 1e5 + 100;
  13. const int N = 1e5 + 10;
  14. const int MOD = 1e9 + 7;
  15.  
  16. array<array<int, MAXN>, 3> dp;
  17. array<int, MAXN> color, used;
  18. array<vector<int>, MAXN> g;
  19.  
  20. template<typename T>
  21. T operator+(const T& a, const T& b) {
  22.     return (a + b) % MOD;
  23. }
  24.  
  25. void paint_it(int v, int to) {
  26.     if (color[v] == -1) {
  27.         dp[v][0] = dp[to][0];
  28.         dp[v][1] = dp[to][1];
  29.         dp[v][2] = dp[to][2];
  30.     } else {
  31.         dp[v][0] = dp[v][0] + dp[to][1] + dp[to][2];
  32.         dp[v][1] = dp[v][1] + dp[to][0] + dp[to][1];
  33.         dp[v][2] = dp[v][2] + dp[to][0] + dp[to][1];
  34.         if (color[to] != -1) {
  35. //            dp[v][color[to]] = 0;
  36.         }
  37.     }
  38. }
  39.  
  40. void dfs(int v, int p) {
  41.     used[v] = true;
  42.     bool flag = true;
  43.     for (auto to: g[v]) {
  44.         if (!used[to]) {
  45.             dfs(to, v);
  46.             flag = false;
  47.         }
  48.         if (to != p) {
  49.             paint_it(v, to);
  50.         }
  51.     }
  52.     if (flag) {
  53.         if (color[v] != -1) {
  54.             dp[v][color[v]] = 1;
  55.         } else {
  56.             dp[v][0] = 1;
  57.             dp[v][1] = 1;
  58.             dp[v][2] = 1;
  59.         }
  60.     }
  61. }
  62.  
  63. signed main() {
  64.     ios_base::sync_with_stdio(false);
  65.     cin.tie(nullptr);
  66.     cout.tie(nullptr);
  67.     int n, k;
  68.     cin >> n >> k;
  69.     for (int i = 0; i < n - 1; ++i) {
  70.         int u, v;
  71.         cin >> u >> v;
  72.         --u, --v;
  73.         g[u].emplace_back(v);
  74.         g[v].emplace_back(u);
  75.     }
  76.     color.fill(-1);
  77.     for (int i = 0; i < k; ++i) {
  78.         int v, col;
  79.         cin >> v >> col;
  80.         --v, --col;
  81.         color[v] = col;
  82.     }
  83.     for (int i = 0; i < n; ++i) {
  84.         for (int j = 0; j < 3; ++j) {
  85.             dp[i][j] = 0;
  86.         }
  87.     }
  88.     dfs(0, -1);
  89.     cout << dp[0][0] << ' ' << dp[0][1] << ' ' << dp[0][2] << endl;
  90.     cout << dp[0][0] + dp[0][1] + dp[0][2] << '\n';
  91. }
Advertisement
Add Comment
Please, Sign In to add comment