$$\text{Obviously,} \dp_u = \forall x, y\in \mathrm{sons}_u \operatorname{and} x\neq y, \\forall a \in \mathrm{subtree}_x, b\in \mathrm{subtree}y, \operatorname{and} a, b\in \mathrm{key}, \\sum \operatorname{dis}(a, b) \+ \sum{v \in \mathrm{sons}u} dp_v + [u\in \mathrm{key}]\sum{v \in \mathrm{subtree}_u, v\in \mathrm{key}} \operatorname{dis}(u, v), \\\text{What's more, we have} \\operatorname{dis}(a, b) = \mathrm{dep}_a + \mathrm{dep}b - 2\mathrm{dep}{\operatorname{lca}(a,b)}, \\text{Therefore, }\\forall a \in \mathrm{subtree}_x, b\in \mathrm{subtree}_y, \operatorname{and} a, b\in \mathrm{key}, \\operatorname{dis}(a, b) = \mathrm{dep}_a + \mathrm{dep}b - 2 \mathrm{dep}{\operatorname{lca}(x, y)} = \mathrm{dep}_a + \mathrm{dep}_b - 2 \mathrm{dep}u, \\text{Let's write it in a $\Sigma$ format: } \\begin{aligned}dp_u & = \sum{v \in \mathrm{sons}u} dp_v + [u\in \mathrm{key}]\sum{v \in \mathrm{subtree}_u, v\in \mathrm{key}} (\mathrm{dep}_v - \mathrm{dep}u) + \sum{x, y\in \mathrm{sons}u, x\neq y}~\sum{a \in \mathrm{subtree}x, a\in \mathrm{key}} ~ \sum{b\in \mathrm{subtree}_y, b\in \mathrm{key}} (\mathrm{dep}_a + \mathrm{dep}_b - 2 \mathrm{dep}u) \ & = \sum{v \in \mathrm{sons}u} dp_v + [u\in \mathrm{key}]\left(\sum{v \in \mathrm{subtree}_u, v\in \mathrm{key}} \mathrm{dep}_v - \mathrm{cnt}_u\mathrm{dep}u\right) + \sum{x, y\in \mathrm{sons}u, x\neq y}~\sum{a \in \mathrm{subtree}_x, a\in \mathrm{key}}\mathrm{dep}_a\mathrm{cnt}_y - 2 \mathrm{dep}_u\mathrm{cnt}y+\sum{b\in \mathrm{subtree}_y, b\in \mathrm{key}} \mathrm{dep}_b \\end{aligned}\\text{Where $\mathrm{cnt}_u$ is the count of key nodes in the subtree of $u$.}\\text{Denote that $\mathrm{sum}u$ is the sum of the depth of key nodes in the subtree of $u$, then we have}\\begin{aligned}dp_u & = \sum{v \in \mathrm{sons}_u} dp_v + [u\in \mathrm{key}]\left(\mathrm{sum}_u - \mathrm{cnt}_u \mathrm{dep}u\right) + \sum{x, y\in \mathrm{sons}u, x\neq y}~\sum{a \in \mathrm{subtree}_x, a\in \mathrm{key}}\mathrm{dep}_a\mathrm{cnt}_y - 2 \mathrm{dep}_u\mathrm{cnt}_y+\mathrm{sum}y \\end{aligned}\\text{But it's still $O(n^3)$, quite slow. Note there are a lot useless $\Sigma$-s, how to remove them?}\\begin{aligned}dp_u & = \sum{v \in \mathrm{sons}_u} dp_v + [u\in \mathrm{key}]\left(\mathrm{sum}_u - \mathrm{cnt}_u \mathrm{dep}u\right) + \sum{x, y\in \mathrm{sons}u, x\neq y}~\sum{a \in \mathrm{subtree}_x, a\in \mathrm{key}}\mathrm{dep}_a\mathrm{cnt}_y - 2 \mathrm{dep}_u\mathrm{cnt}_y+\mathrm{sum}y \ & = \sum{v \in \mathrm{sons}_u} dp_v + [u\in \mathrm{key}]\left(\mathrm{sum}_u - \mathrm{cnt}_u \mathrm{dep}u\right) + \sum{x, y\in \mathrm{sons}_u, x\neq y} \mathrm{cnt}_x(- 2 \mathrm{dep}_u\mathrm{cnt}_y+\mathrm{sum}y) + \sum{a \in \mathrm{subtree}_x, a\in \mathrm{key}}\mathrm{dep}_a\mathrm{cnt}y \ & = \sum{v \in \mathrm{sons}_u} dp_v + [u\in \mathrm{key}]\left(\mathrm{sum}_u - \mathrm{cnt}_u \mathrm{dep}u\right) + \sum{x, y\in \mathrm{sons}_u, x\neq y}\mathrm{sum}_y\mathrm{cnt}_x + \mathrm{sum}_x\mathrm{cnt}_y - 2 \mathrm{dep}_u\mathrm{cnt}_x\mathrm{cnt}y \\end{aligned}\\text{The first two parts can be easily solved in $O(n)$ time complexity. To solve the third part, we could rearrange it: }\\begin{aligned}dp_u & = \sum{v \in \mathrm{sons}_u} dp_v + [u\in \mathrm{key}]\left(\mathrm{sum}_u - \mathrm{cnt}_u \mathrm{dep}u\right) + \sum{x, y\in \mathrm{sons}_u, x\neq y}\mathrm{sum}_y\mathrm{cnt}_x + \mathrm{sum}_x\mathrm{cnt}_y - 2 \mathrm{dep}_u\mathrm{cnt}_x\mathrm{cnt}y \ & = \sum{v \in \mathrm{sons}_u} dp_v + [u\in \mathrm{key}]\left(\mathrm{sum}_u - \mathrm{cnt}_u \mathrm{dep}u\right) + \sum{x \in \mathrm{sons}u}\sum{y \in \mathrm{sons}_u} [x\neq y]~ \mathrm{sum}_x\mathrm{cnt}_y - \mathrm{dep}_u\mathrm{cnt}_x\mathrm{cnt}_y \\end{aligned}\\text{It seems to be much more clean. }\\text{Note we removed the $2$ before $\mathrm{dep}_u\mathrm{cnt}_x\mathrm{cnt}y$,}\\text{that's because for each unordered pair $(x, y)$, we calculated it twice as ordered pair $(x, y), (y, x).$}\\text{To move on, we should first remove $[x\neq y]$, which is annoying. We can just ignore it, and remove it's contribution afterwards.}\\begin{aligned}dp_u & = \sum{v \in \mathrm{sons}_u} dp_v + [u\in \mathrm{key}]\left(\mathrm{sum}_u - \mathrm{cnt}_u \mathrm{dep}u\right) + \sum{x \in \mathrm{sons}u}\sum{y \in \mathrm{sons}_u} \mathrm{sum}_x\mathrm{cnt}_y - \mathrm{dep}_u\mathrm{cnt}_x\mathrm{cnt}y - \sum{x \in \mathrm{sons}_u} \mathrm{sum}_x\mathrm{cnt}_x - \mathrm{dep}_u\mathrm{cnt}x^2 \ & = \sum{v \in \mathrm{sons}_u} dp_v + [u\in \mathrm{key}]\left(\mathrm{sum}_u - \mathrm{cnt}_u \mathrm{dep}u\right) + \sum{x \in \mathrm{sons}u}\sum{y \in \mathrm{sons}_u} (\mathrm{sum}_x - \mathrm{dep}_u\mathrm{cnt}_x)\mathrm{cnt}y - \sum{x \in \mathrm{sons}_u} \mathrm{sum}_x\mathrm{cnt}_x - \mathrm{dep}_u\mathrm{cnt}x^2 \ & = \sum{v \in \mathrm{sons}_u} dp_v + [u\in \mathrm{key}]\left(\mathrm{sum}_u - \mathrm{cnt}_u \mathrm{dep}u\right) + \sum{x \in \mathrm{sons}_u}(\mathrm{sum}_x - \mathrm{dep}_u\mathrm{cnt}x)\sum{y \in \mathrm{sons}_u} \mathrm{cnt}y - \sum{x \in \mathrm{sons}_u} \mathrm{sum}_x\mathrm{cnt}_x - \mathrm{dep}_u\mathrm{cnt}_x^2 \\end{aligned}\\text{These parts are all independent with each other. This means, we finally found a $O(n)$ solution!}\\tiny\text{Though I don't know if it's right or not \dots}\normalsize\\text{The second and third subtasks are rather easy. }\\text{For example, if we want to calculate the longest distance, }\\text{First define $\mathrm{deepest}u$ as the deepest key node in the subtree of $u$, then: }\dp_u = \max\left( \max\limits{v\in \mathrm{sons}u}dp_v, [u \in \mathrm{key}] ~ \mathrm{dep}{\mathrm{deepest}_u} - \mathrm{dep}u, \max\limits{v\in \mathrm{sons}u}\max\limits{w\in \mathrm{sons}_u}\operatorname{dis}(\mathrm{deepest}_v,\mathrm{deepest}_w) \right)\\text{The last part can be solved by recording the deepest node so far.}\\text{Note: to calculate the shortest distance, more edge cases must be tested.}\$$