Kitty's Calculations on a Tree

                            Solution in C :

In C ++ :




#include <bits/stdc++.h>

using namespace std;

const int MOD=1000000007;
int N, Q;
vector<int> adj[200001];
vector<int> adj2[200001];
int P[18][200001];
int depth[200001];
int in[200001];
int out[200001];
int now;
int A[400001];
int mul[200001];
long long sum[200001];

void dfs(int u, int p)
{
    P[0][u]=p;
    for(int i=1; i<18; i++)
        P[i][u]=P[i-1][P[i-1][u]];
    in[u]=++now;
    for(auto& v: adj[u]) if(v!=p)
    {
        depth[v]=depth[u]+1;
        dfs(v, u);
    }
    out[u]=now;
}

int lca(int u, int v)
{
    if(depth[u]<depth[v])
        swap(u, v);
    for(int i=17; i>=0; i--) if(depth[P[i][u]]>=depth[v])
        u=P[i][u];
    if(u==v)
        return u;
    for(int i=17; i>=0; i--) if(P[i][u]!=P[i][v])
        u=P[i][u], v=P[i][v];
    return P[0][u];
}

int dfs2(int u, long long tot)
{
    int ret=0;
    sum[u]=u*mul[u];
    for(auto& v: adj2[u])
    {
        ret=(ret+dfs2(v, tot))%MOD;
        sum[u]+=sum[v];
    }
    for(auto& v: adj2[u])
        ret=(ret+1LL*((tot-sum[v])%MOD)
                    *(sum[v]%MOD)%MOD
                    *(depth[v]-depth[u])%MOD)%MOD;
    return ret;
}

int main()
{
    scanf("%d%d", &N, &Q);
    for(int i=0; i<N-1; i++)
    {
        int a, b;
        scanf("%d%d", &a, &b);
        adj[a].push_back(b);
        adj[b].push_back(a);
    }
    dfs(1, 1);
    while(Q--)
    {
        int K;
        scanf("%d", &K);
        long long tot=0;
        for(int i=0; i<K; i++)
            scanf("%d", A+i), mul[A[i]]=1, tot+=A[i];
        sort(A, A+K, [](int a, int b) {
            return in[a]<in[b];
        });
        for(int i=0; i<K-1; i++)
            A[i+K]=lca(A[i], A[i+1]);
        sort(A, A+2*K-1);
        int M=unique(A, A+2*K-1)-A;
        sort(A, A+M, [](int a, int b) {
            return out[a]-in[a]>out[b]-in[b];
        });
        int root=A[0];
        map<int, int> m;
        m[in[root]]=root;
        for(int i=1; i<M; i++)
        {
            int u=A[i];
            auto it=m.upper_bound(in[u]);
            assert(it!=m.begin());
            --it;
            int p=it->second;
            adj2[p].push_back(u);
            //printf("%d -> %d\n", p, u);
            m[in[u]]=u;
            if(out[u]<out[p] && (!m.count(out[u]+1) || P[0][m[out[u]+1]]!=p))
                m[out[u]+1]=p;
        }
        printf("%d\n", dfs2(root, tot));
        for(int i=0; i<M; i++)
            adj2[A[i]].clear(), mul[A[i]]=0;
    }
    return 0;
}








In Java :




import java.util.Arrays;
import java.util.Scanner;

public class KittysCalc {
	
	public static final long constant = 1000000007;
	
	public static void main(String[] args) {
		
	Scanner sc = new Scanner(System.in);
	int n = sc.nextInt();
	int queries = sc.nextInt();
	int[] parents = new int[n+1];
	long[] children = new long[n+1];
	boolean[] valuesSet = new boolean[n+1];
	long valuesSum = 0;
	long sum = 0;
	int a, b;
	for(int i = 0; i < n-1; i++) {
		a = sc.nextInt();
		b = sc.nextInt();
		if(a < b) {
			parents[b] = a;
		} else {
			parents[a] = b;
		}
	}
	parents[1] = 0;
		
	for(int i = 0; i < queries; i++) {
		int k = sc.nextInt();
		Arrays.fill(valuesSet, false);
		Arrays.fill(children, 0);
		valuesSum = 0;
		for(int j = 0; j < k; j++) {
			a = sc.nextInt();
			valuesSum += a;
			valuesSet[a] = true;
		}
		sum = 0;
		for (int j = n; j > 0; j--) {
		long c = children[j];
		if (valuesSet[j]) {
			c += j;
		}
		if (c > 0) {
		long x = ((c % constant) * ((valuesSum - c) % constant)) % constant;
		if (constant - sum < x) {
			sum -= constant;
		}
		sum += x;
			}
	children[parents[j]] += c;
			}
System.out.println(sum);
		}
		sc.close();
	}

}








In Python3 :





#!/usr/bin/env python3

def put(d, a, b):
    if a in d: d[a].append(b)
    else: d[a] = [b]

def main():
    for n in ns[::-1]:
        r = [tt[s] for s in tree[n] if s != f[n]]
        bst = {s: [gl[n], n, 0] for s in queries[n]}
        if r:
            o = max(range(len(r)), key=lambda a: len(r[a]))
            if len(r[o]) > len(bst): r[o], bst = bst, r[o]
        ry = {}
        for ae in r:
            for y, v in ae.items():
                put(ry, y, v)
        for y, r in ry.items():
            eq, z, t = 0, 0, 0
            if len(r) == 1 and y not in bst:
                bst[y] = r[0]
                continue
            if y in bst: r.append(bst.pop(y))
            for d, v, c in r:
                eq += (d - gl[n]) * v + c
                z += v
            for d, v, c in r:
                c += (d - gl[n]) * v
                diff = (eq - c) * v
                t += diff
            returns[y] += t
            bst[y] = (gl[n], z, eq)
        tt[n] = bst
    
def locate():
    q = [r]
    level = 0
    while q:
        level += 1
        tmp = []
        ns.extend(q)
        for n in q:
            for s in tree[n]:
                if s not in f:
                    f[s] = n
                    gl[s] = level
                    tmp.append(s)
        q = tmp
        
tree = {}
tt = {}
n, q = map(int, input().split())
returns = [0] * q
for _ in range(n - 1):
    a, b = map(int, input().split())
    put(tree, a, b)
    put(tree, b, a)
queries = {a: set() for a in tree}
for y in range(q):
    input()
    for x in map(int, input().split()): queries[x].add(y)
r = next(iter(tree))
ns = []
f = {r: None}
gl = {r: 0}
locate()
main()
for s in returns: print(s % (10**9 + 7))