Kitty's Calculations on a Tree

Problem Statement :

Kitty has a tree, T , consisting of n nodes where each node is uniquely labeled from  1 to n . Her friend Alex gave her q sets, where each set contains k distinct nodes. Kitty needs to calculate the following expression on each set:


{ u ,v } denotes an unordered pair of nodes belonging to the set.
 dist(u , v) denotes the number of edges on the unique (shortest) path between nodes  and .
Given T and q sets of k  distinct nodes, calculate the expression for each set. For each set of nodes, print the value of the expression modulo 10^9 + 7  on a new line.

Input Format

The first line contains two space-separated integers, the respective values of n (the number of nodes in tree T ) and  q (the number of nodes in the query set).
Each of the n - 1  subsequent lines contains two space-separated integers, a and b, that describe an undirected edge between nodes  and .
The 2 * q subsequent lines define each set over two lines in the following format:

The first line contains an integer, k  , the size of the set.
The second line contains  k space-separated integers, the set's elements.

Output Format

Print q lines of output where each line i contains the expression for the ith query, modulo 10^9 + 7.

Solution :


                            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)
    for(int i=1; i<18; i++)
    for(auto& v: adj[u]) if(v!=p)
        dfs(v, u);

int lca(int u, int v)
        swap(u, v);
    for(int i=17; i>=0; i--) if(depth[P[i][u]]>=depth[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;
    for(auto& v: adj2[u])
        ret=(ret+dfs2(v, tot))%MOD;
    for(auto& v: adj2[u])
    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);
    dfs(1, 1);
        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;
        for(int i=1; i<M; i++)
            int u=A[i];
            auto it=m.upper_bound(in[u]);
            int p=it->second;
            //printf("%d -> %d\n", p, u);
            if(out[u]<out[p] && (!m.count(out[u]+1) || P[0][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(;
	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;


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]
            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 = []
        for n in q:
            for s in tree[n]:
                if s not in f:
                    f[s] = n
                    gl[s] = level
        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):
    for x in map(int, input().split()): queries[x].add(y)
r = next(iter(tree))
ns = []
f = {r: None}
gl = {r: 0}
for s in returns: print(s % (10**9 + 7))

View More Similar Problems


An array is a type of data structure that stores elements of the same type in a contiguous block of memory. In an array, A, of size N, each memory location has some unique index, i (where 0<=i<N), that can be referenced as A[i] or Ai. Reverse an array of integers. Note: If you've already solved our C++ domain's Arrays Introduction challenge, you may want to skip this. Example: A=[1,2,3

View Solution →

2D Array-DS

Given a 6*6 2D Array, arr: 1 1 1 0 0 0 0 1 0 0 0 0 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 An hourglass in A is a subset of values with indices falling in this pattern in arr's graphical representation: a b c d e f g There are 16 hourglasses in arr. An hourglass sum is the sum of an hourglass' values. Calculate the hourglass sum for every hourglass in arr, then print t

View Solution →

Dynamic Array

Create a list, seqList, of n empty sequences, where each sequence is indexed from 0 to n-1. The elements within each of the n sequences also use 0-indexing. Create an integer, lastAnswer, and initialize it to 0. There are 2 types of queries that can be performed on the list of sequences: 1. Query: 1 x y a. Find the sequence, seq, at index ((x xor lastAnswer)%n) in seqList.

View Solution →

Left Rotation

A left rotation operation on an array of size n shifts each of the array's elements 1 unit to the left. Given an integer, d, rotate the array that many steps left and return the result. Example: d=2 arr=[1,2,3,4,5] After 2 rotations, arr'=[3,4,5,1,2]. Function Description: Complete the rotateLeft function in the editor below. rotateLeft has the following parameters: 1. int d

View Solution →

Sparse Arrays

There is a collection of input strings and a collection of query strings. For each query string, determine how many times it occurs in the list of input strings. Return an array of the results. Example: strings=['ab', 'ab', 'abc'] queries=['ab', 'abc', 'bc'] There are instances of 'ab', 1 of 'abc' and 0 of 'bc'. For each query, add an element to the return array, results=[2,1,0]. Fun

View Solution →

Array Manipulation

Starting with a 1-indexed array of zeros and a list of operations, for each operation add a value to each of the array element between two given indices, inclusive. Once all operations have been performed, return the maximum value in the array. Example: n=10 queries=[[1,5,3], [4,8,7], [6,9,1]] Queries are interpreted as follows: a b k 1 5 3 4 8 7 6 9 1 Add the valu

View Solution →