fork download
copy 
  1. #include <iostream>
  2. #include <vector>
  3. using namespace std;
  4.  
  5. const int MOD = 998244353;
  6.  
  7. int main() {
  8.     ios::sync_with_stdio(false);
  9.     cin.tie(nullptr);
  10.  
  11.     int t;
  12.     cin >> t;
  13.     while (t--) {
  14.         int n;
  15.         cin >> n;
  16.  
  17.         // Allocate array for depth; vertices are numbered 1..n.
  18.         vector<int> depth(n + 1, 0); // root has depth 0
  19.         int max_depth = 0;
  20.  
  21.         // Read parent for vertices 2..n and compute depth.
  22.         for (int i = 2; i <= n; i++) {
  23.             int p;
  24.             cin >> p;
  25.             depth[i] = depth[p] + 1;
  26.             if (depth[i] > max_depth) {
  27.                 max_depth = depth[i];
  28.             }
  29.         }
  30.  
  31.         // Count vertices by depth (for depths >= 1).
  32.         vector<long long> cnt(max_depth + 1, 0);
  33.         for (int i = 2; i <= n; i++) {
  34.             cnt[depth[i]]++;
  35.         }
  36.  
  37.         // S[d] = sum of f(v) for all v at depth d, for d >= 1.
  38.         vector<long long> S(max_depth + 2, 0);
  39.         // Base case: deepest level
  40.         S[max_depth] = cnt[max_depth] % MOD;
  41.  
  42.         // Use recurrence: S[d] = cnt[d] + (cnt[d] - 1) * S[d+1]
  43.         for (int d = max_depth - 1; d >= 1; d--) {
  44.             long long ways = cnt[d] % MOD;
  45.             long long factor = (cnt[d] - 1) % MOD;
  46.             if (factor < 0) {
  47.                 factor += MOD; // not needed as cnt[d] >= 1 in a tree
  48.             }
  49.             S[d] = (ways + (factor * S[d + 1]) % MOD) % MOD;
  50.         }
  51.  
  52.         // f(1) = 1 (the sequence with just root) plus sequences from moves from the root.
  53.         long long result = (1 + S[1]) % MOD;
  54.         cout << result << "\n";
  55.     }
  56.  
  57.     return 0;
  58. }
Success #stdin #stdout 0.01s 5284KB
comments (?)
stdin
copy 
3
4
1 2 1
3
1 2
7
1 2 2 1 4 5
stdout
copy 
4
2
8
https://ideone.com/XdFFhQ
language:
C++ (gcc 8.3)
created:
10 days 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!