Beautiful Tree — An interview question

Hey guys !!

Hope you all are in good condition and enjoying the lock down ;)

I recently stumbled on a question of trees and online queries and I'm unable to solve it. So I'm attaching the image containing the question. And, just in case if the image is not clear, I'm writing it in the same manner (or at least trying to) as it is in the attached image. And please don't ask for the link to the question as I don't have one.

You are given a tree with N vertices, rooted at vertex 1 (indexed from 1 to N). Each vertex has a value associated with it. The value of $$$i_{th}$$$ vertex is $$$a_i$$$. Consider $$$f_i$$$ as the $$$a_i$$$ th beautiful number.

A beautiful number is a number whose sum of squares of digits is less than or equal to X. You have to process M queries where each query can be of two types.

Type 1 -> 1 i y : update $$$a-i$$$ to y

Type 2 -> 2 i : output the sum of $$$f_i$$$ for all the nodes in the sub-tree of node i ( including node i). Since the number can be very large compute it modulo 998244353.

CONSTRAINTS :

1 <= N, M, i, $$$u_i$$$, $$$v_i$$$, $$$a_i$$$, y <= 5*$$$10^5$$$

1 <= X <= $$$10^{10}$$$

Time Limit : 1 sec

If anyone got any ideas or know how to solve this questions in the given constraint, please do tell.

Also under what categories such questions fall ?

And it would be great help if anyone can share a bunch of resources from where beginners like me can learn and practice these type of questions.

Thank you.