Beautiful Tree

#include <bits/stdc++.h>
#define int long long

using namespace std;

const int MAX=500000;
const int MOD=998244353;

vector<int> beautiful;
set<int> st;
int cnt;
int n,m,x;

vector<vector<int> > graph;  // adjacency list of graph
vector<bool> vis;  // visited array
vector<int> in;  // discover times
vector<int> out;  // finish times
vector<int> nodes; // nodes in order of their discovery time
int tot;  // total nodes visited yet

void dfs(int i)
{
   vis[i]=true;
   nodes.push_back(i);
   in[i]=tot++;
   for(auto it: graph[i])
    {
        if(!vis[it])
            dfs(it);
    }
   out[i]=tot-1;
}

 class BIT
{
    // 1-based indexing
    vector<int> bit;
    int n;
    public:
    BIT(int sz)
    {
        n=sz;
        bit.assign(n+1,0);
    }
    void update(int x, int del)
    {
        // queries are 1-based
        while(x<=n)
        {
            bit[x]=(((bit[x]+(del%MOD))%MOD)+MOD)%MOD;
            x+=x&(-x);
        }
    }
    int query(int x)
    {
        int ans=0;
        while(x>0)
        {
            ans=(ans+bit[x])%MOD;
            ans=(ans+MOD)%MOD;
            x-=x&(-x);
        }
        return ans;
    }
    int query(int l, int r)
    {
        return (((query(r)-query(l-1))%MOD)+MOD)%MOD;
    }
};

int32_t main()
{
  cin>>n>>m>>x;
  cnt=0;
  queue<pair<int,int> > q;

  for(int i=1; i<=9; i++)
  {
    if(i*i<=x)
    {
        if(cnt+1>MAX)
            break;
        cnt++;
        beautiful.push_back(i);
        q.push({i,i*i});
    }
    else
        break;
  }

  while(!q.empty())
  {
     if(cnt+1>MAX)
        break;
    int  num=q.front().first;
    int sum=q.front().second;
      q.pop();
      for(int i=0; i<=9; i++)
      {
        if(sum+i*i<=x)
        {
            int newNum=(((num*10)%MOD)+i)%MOD;
            if(cnt+1>MAX)
                break;
            cnt++;
            beautiful.push_back(newNum);
            q.push({newNum,sum+i*i});
        }
        else
            break;
      }
  }
  graph.resize(n+1);
  for(int i=0; i<n-1; i++)
  {
      int u,v;
      cin>>u>>v;
      graph[u].push_back(v);
      graph[v].push_back(u);
  }

  vis.assign(n+1,false);
  in.resize(n+1);
  out.resize(n+1);
  tot=1;
  nodes.push_back(0); //dummy element, since i kept 1-based indexing
  dfs(1);
  vector<int> val(n+1);
  BIT b(n);
  for(int i=1; i<=n; ++i)
  {
      cin>>val[i];
      int ind=val[i]-1;
      b.update(in[i],beautiful[ind]);
  }
  while(m--)
  {
      int type;
      cin>>type;
      if(type==1)
      {
          int i,y;
          cin>>i>>y;
          int prev=beautiful[val[i]-1];
          int curr=beautiful[y-1];
          int del=curr-prev;
          b.update(in[i],del);
          val[i]=y;
      }
      else
      {
          int i;
          cin>>i;
          int l=in[i];
          int r=out[i];
          cout<<b.query(l,r)<<'\n';
      }
  }
  return 0;
}

/*  Sample 1:
3 3 5
1 2
1 3
1 2 3
2 1
1 2 5
2 1
* */

/*  Sample 2:
3 3 5
1 2
1 3
1 2 3
2 1
1 2 98
2 1
* */