Basecamp
    Home
    Problems
    Contests
    Courses
    Rating
    Posts
    Discord
Trees
Sign in
medv
•Article•1 month ago

Trees

Trees are one of the key data structures in algorithms and appear in a wide range of problems.

A tree is a connected, acyclic, undirected graph.

A graph G=(V,E) is called a tree if it satisfies the following properties:

  • Connected — there is a path between every pair of vertices in the graph.

  • Acyclic — the graph contains no cycles.

In a tree with n vertices, there are always exactly n−1 edges.

Fundamental Properties of a Tree

1. Connectivity and Uniqueness of Paths

In a tree, there is a unique path between any pair of vertices.

2. Number of Edges

For any tree:

∣E∣=∣V∣−1

3. No Cycles

Adding even a single edge to a tree always creates a cycle.

4. Root (for a rooted tree)

If one vertex is chosen as the root, the tree becomes directed from the root to its descendants. This is used in recursive traversals and algorithms (DFS, LCA, etc.).

5. Subtrees

Any vertex along with its descendants forms a subtree. A subtree is itself a tree.

6. Leaves

Vertices with zero children (in a rooted tree) are called leaves.

7. Depth and Height

  • Depth of a vertex is the distance from the root to that vertex.

  • Height of a tree is the maximum depth among all vertices.

8. Center of a Tree

The center is the vertex (or a pair of adjacent vertices) that minimizes the maximum distance to all other vertices.

Alternative Definitions of a Tree

The following definitions of a tree are equivalent and can be used in different contexts:

  • A cycle-free graph that becomes disconnected if any edge is removed.

  • A connected graph with n vertices and n−1 edges.

  • An acyclic graph with n−1 edges and exactly one connected component.

Representing Trees in Code

Depending on the problem and the type of tree (general, binary, rooted, etc.), different data structures are used:

1. Adjacency Lists

This approach is often used in competitive programming problems and when working with undirected or directed trees.

int n;
vector<vector<int>> tree(n); // tree with n vertices

// adding an edge between vertices u and v
tree[u].push_back(v);
tree[v].push_back(u); // for an undirected tree

2. Parent Array

A very compact representation, especially useful when the tree is already constructed.

vector<int> parent(n); // parent[i] --- he parent of vertex i
// if i is the root, parent[i] is usually set to -1 or 0

3. Pointers / Structures (OOP Style)

Commonly used when constructing binary trees, tries, decorators, especially in Python/Java/C++ with object-oriented programming.

  • Binary Tree:

struct Node 
{
  int val;
  Node* left;
  Node* right;
  Node(int v) : val(v), left(nullptr), right(nullptr) {}
};
C++
7 lines
122 bytes
  • General Tree:

struct Node 
{
  int val;
  vector<Node*> children;
  Node(int v) : val(v) {}
};
Eolymp #977. Is it a Tree?

An undirected graph without loops and multiple edges is given by an adjacency matrix. Determine whether this graph is a tree.

Input. The first line contains the number of vertices n (1≤n≤100) in the graph. Then, an adjacency matrix of size n×n is given, where 1 indicates the presence of an edge, and 0 indicates its absence. The matrix is symmetric with respect to the main diagonal.

Output. Print "YES" if the graph is a tree, and "NO" otherwise.

Sample input
3
0 1 0
1 0 1
0 1 0
Sample output
YES
Open problem
Solution

A graph with n vertices is a tree if and only if:

  1. The graph is connected. Start a depth-first search from the first vertex. Count the number of visited vertices during the search. If it equals n, then the graph is connected.

  2. ∣V∣=∣E∣+1. If the graph is a tree, then it contains n−1 edges

  3. The graph does not contain cycles. Start a depth-first search from the first vertex. If a back edge exists, then the graph has a cycle and is not a tree.

Satisfying conditions 1) and 2) or 1) and 3) is sufficient for the graph to be a tree.

Example

The graph provided in the example is a tree.

Algorithm implementation --- checking the conditions 1) and 3)

Declare the adjacency matrix of the graph g and the array used.

#define MAX 110
int g[MAX][MAX], used[MAX];

The dfs function performs a depth-first search starting from vertex v. The parent of vertex v is p. If a cycle is detected, the flag variable is set to 1.

void dfs(int v, int p)
{

If the graph is no longer a tree (flag=1), there is no need to continue the search.

  if (flag) return;

Mark the vertex v as visited.

  used[v] = 1;

The vertex v is visited, increase c by 1.

  c++;

The edge (v,i) will be a back edge and form a cycle if i=p and vertex i is already visited (used[i]=1). If a cycle is detected, set flag=1. If no cycle is detected, continue the search from vertex i.

  for(int i = 0; i < n; i++)
    if ((i != p) && g[v][i])
      if(used[i]) flag = 1; else dfs(i,v);
}

The main part of the program. Read the input data.

scanf("%d",&n); 
for(i = 0; i < n; i++)
for(j = 0; j < n; j++)
  scanf("%d",&g[i][j]);

Count the number of visited vertices during the depth-first search in the variable c. Set flag=0 if there is no cycle in the graph. If a cycle is detected, flag becomes 1.

c = 0;
flag = 0;

All vertices are initially unvisited (initialize the array used with zeroes).

memset(used,0,sizeof(used));

Start the depth-first search from vertex 0. Since it is the root of the tree, it has no parent. Pass the value −1 as the second argument to the dfs function.

dfs(0,-1);

The graph is not a tree if there is a back edge (flag=1) or if the graph is not connected (c=n).

if (flag || (c != n)) printf("NO\n"); else printf("YES\n");

Python implementation — checking the conditions 1) and 3)

The dfs function performs a depth-first search starting from vertex v. The parent of vertex v is p. If a cycle is detected, the flag variable is set to 1.

def dfs(v, p):

Declare the global variables used by the function.

  global flag, c

If the graph is no longer a tree (flag=1), there is no need to continue the search.

  if flag: return

Mark the vertex v as visited.

  used[v] = 1

The vertex v is visited, increase c by 1.

  c += 1

The edge (v,i) will be a back edge and form a cycle if i=p and vertex i is already visited (used[i]=1). If a cycle is detected, set flag=1. If no cycle is detected, continue the search from vertex i.

  for i in range(n):
    if i != p and g[v][i]:
      if used[i]: flag = 1
      else: dfs(i, v)

The main part of the program. Read the input value of n.

n = int(input())

Count the number of visited vertices during the depth-first search in the variable c. Set flag=0 if there is no cycle in the graph. When a cycle is detected, flag becomes 1.

c = 0
flag = 0

All vertices are initially unvisited (initialize the list used with zeroes).

used = [0] * n

Create an adjacency matrix g, initially filled with zeroes.

g = [[0] * n for _ in range(n)]

Read the input adjacency matrix.

for i in range(n):
  g[i] = list(map(int, input().split()))

Start the depth-first search from vertex 0. Since it is the root of the tree, it has no parent. Pass the value −1 as the second argument to the dfs function.

dfs(0, -1)

The graph is not a tree if there is a back edge (flag=1) or if the graph is not connected (c=n).

if flag or c != n:
  print("NO")
else:
  print("YES")
Algorithm implementation --- checking the conditions 1) and 2)

Declare the adjacency matrix of the graph g and the array used.

#define MAX 101
int g[MAX][MAX], used[MAX];

The dfs function implements depth-first search from vertex v.

void dfs(int v)
{

Mark the vertex v as visited.

  used[v] = 1;

Vertex v is visited, increase c by 1.

  c++;

Iterate over the vertices i, that can be reached from v. Moving from v to i is possible if:

  1. There is an edge (v,i), that is g[v][i]=1;

  2. The vertex i is not visited (used[i]=0)

If both conditions are true, we continue the depth-first search from vertex i.

  for (int i = 1; i <= n; i++)
    if (g[v][i] && !used[i]) dfs(i);
}

The main part of the program. Read the input value of n.

scanf("%d", &n); 

Count the number of edges in the graph in the variable Edges.

Edges = 0;

All vertices are initially unvisited (initialize the array used with zeroes).

memset(used, 0, sizeof(used));

Read the adjacency matrix g. Count the number of edges in the graph.

for (i = 1; i <= n; i++)
for (j = 1; j <= n; j++)
{
  scanf("%d", &g[i][j]);
  Edges += g[i][j];
}

Since the graph is undirected, edges (u,v) and (v,u) are considered the same. Divide the value of Edges by 2.

Edges /= 2;

Count the number of visited vertices during the depth-first search in variable c.

c = 0;

Start the depth-first search from the vertex 1.

dfs(1);

The graph is a tree if the number of edges in it equals n−1, and also if it is connected (c=n).

if ((Edges == n - 1) && (c == n)) printf("YES\n");
else printf("NO\n");

Python implementation — checking the conditions 1) and 2)

The dfs function implements depth-first search from vertex v.

def dfs(v):
  global c

Mark the vertex v as visited.

  used[v] = 1

The vertex v is visited, increase c by 1.

  c += 1

Iterate over the vertices i, that can be reached from v. Moving from v to i is possible if:

  1. There is an edge (v,i), that is g[v][i]=1;

  2. The vertex i is not visited (used[i]=0)

If both conditions are true, we continue the depth-first search from vertex i.

  for i in range(1, n + 1):
    if g[v][i] and not used[i]: dfs(i)

The main part of the program. Read the input value of n.

n = int(input())

Count the number of edges in the graph in the variable Edges.

Edges = 0

All vertices are initially unvisited (initialize the array used with zeroes).

used = [0] * (n + 1)

Create an adjacency matrix g, initially filled with zeroes.

g = [[0] * (n + 1) for _ in range(n + 1)]

Read the adjacency matrix g. Count the number of edges in the graph.

for i in range(1, n + 1):
  g[i] = [0] + list(map(int, input().split()))
  Edges += sum(g[i])

Since the graph is undirected, edges (u,v) and (v,u) are considered the same. Divide the value of Edges by 2.

Edges //= 2

Count the number of visited vertices during the depth-first search in variable c.

c = 0

Start the depth-first search from the vertex 1.

dfs(1)

The graph is a tree if the number of edges in it equals n−1, and also if it is connected (c=n).

if Edges == n - 1 and c == n:
  print("YES")
else:
  print("NO")
Eolymp #978.Get a tree

A connected undirected graph without loops or multiple edges is given. It is allowed to delete edges from the graph. The goal is to obtain a tree.

Input. The first line contains two integers – the number of vertices n (1≤n≤100) and the number of edges m in the graph. The following m pairs of integers each represent an edge. It is guaranteed that the graph is connected.

Output. Print n−1 pairs of integers — the edges that form a tree. The edges can be printed in any order.

Sample input
4 4
1 2
2 3
3 4
4 1
Sample output
1 2
2 3
3 4
Open problem
Solution

Perform a depth-first search starting from the first vertex. Build the depth-first search tree and print all of its edges.

Example

The graph given in the example has the following structure:

When the depth-first search is started from vertex 1, the traversal will go through the edges: (1,2),(2,3),(3,4).

Algorithm implementation

The adjacency matrix of the graph is stored in the array g.

#define MAX 101
int g[MAX][MAX], used[MAX];

The dfs function performs a depth-first search starting from vertex v. Each edge (v,i) that is part of the depth-first search tree is printed.

void dfs(int v)
{
  used[v] = 1;
  for (int i = 1; i <= n; i++)
    if (g[v][i] && !used[i])
    {
      printf("%d %d\n",v,i);
      dfs(i);
    }
}
C++
10 lines
160 bytes

The main part of the program. Read the input data.

scanf("%d %d",&n,&m);
memset(g,0,sizeof(g)); memset(used,0,sizeof(used));
while(m--)
{
  scanf("%d %d",&a,&b);
  g[a][b] = g[b][a] = 1;
}
C++
7 lines
145 bytes

Start a depth-first search from the first vertex.

dfs(1);
Eolymp #11263. Even Tree

You are given a tree, which is a simple connected graph without cycles.

Find the maximum number of edges that can be removed from the tree to obtain a forest where each connected component contains an even number of nodes.

For example, in the tree with 4 nodes shown below, you can remove at most 1 edge to create an even forest.

Input. The first line contains two integers n (2≤n≤100,n is even) and m — the number of nodes and edges. Each of the next m lines contains two integers representing the nodes connected by an edge in the tree. The root of the tree is node 1.

Output. Print the maximum number of edges that can be removed.

Sample input
10 9
2 1
3 1
4 3
5 2
6 1
7 2
8 6
9 8
10 8
10 lines
52 bytes
Sample output
2
Open problem
Solution

The task is to decompose the tree into connected components with an even number of vertices, while maximizing the number of such components.

We'll perform a depth-first search starting from vertex 1, which is the root of the tree. For each vertex v, we'll determine the number of vertices in the subtree rooted at v. If this number is even, we'll remove the edge connecting v to its parent in the depth-first search.

For the root (vertex 1), the size of the entire tree is even. However, there is no edge connecting vertex 1 to its parent, so we must subtract 1 from the result obtained using the above algorithm.

Example

Consider the tree shown in the example. Next to each vertex is the size of the subtree when that vertex is considered as the root. Vertices 1,3, and 6 have subtrees with an even number of vertices. Since vertex 1 is the root and does not have an edge leading to a parent, we remove two edges: (1,3) and (1,6).

The tree has been decomposed into 3 connected components, each containing an even number of vertices.

Algorithm implementation

Declare the arrays.

#define MAX 101
int g[MAX][MAX], used[MAX];

The dfs function performs a depth-first search starting from vertex v and returns the number of vertices in the subtree rooted at v.

int dfs(int v)
{
  used[v] = 1;

In the variable s, compute the number of vertices in the subtree rooted at v.

  int s = 1;

Iterate over the children i of vertex v. Add the values of dfs(i) to s.

  for (int i = 1; i <= n; i++)
    if (g[v][i] && !used[i]) s += dfs(i);

If the size of the subtree s is even, remove the edge between v and its parent.

  if (s % 2 == 0) res++;

Return the number of vertices in the subtree rooted at v.

  return s;
}

The main part of the program. Read the input data and build the graph.

scanf("%d %d", &n, &m);
for (i = 0; i < m; i++)
{
  scanf("%d %d", &a, &b);
  g[a][b] = g[b][a] = 1;
}

Start a depth-first search from vertex 1. In the variable res, count the number of removed edges.

res = 0;
dfs(1);

Since there is no edge from 1 to its parent, print res−1.

printf("%d\n", res - 1);

Python implementation

The dfs function performs a depth-first search starting from vertex v and returns the number of vertices in the subtree rooted at v.

def dfs(v):
  used[v] = 1

In the variable s, compute the number of vertices in the subtree rooted at v.

  s = 1

Iterate over the children i of vertex v. Add the values of dfs(i) to s.

  for i in range(1, n + 1):
    if g[v][i] and not used[i]:
      s += dfs(i)

If the size of the subtree s is even, remove the edge between v and its parent.

  if s % 2 == 0:
    global res
    res += 1

Return the number of vertices in the subtree rooted at v.

  return s

The main part of the program. Read the input data and build the graph.

n, m = map(int, input().split())
g = [[0] * (n + 1) for _ in range(n + 1)]
used = [0] * (n + 1)

for _ in range(m):
  a, b = map(int, input().split())
  g[a][b] = g[b][a] = 1
Python
7 lines
182 bytes

Start a depth-first search from vertex 1. In the variable res, count the number of removed edges.

res = 0
dfs(1)

Since there is no edge from 1 to its parent, print res−1.

print(res - 1)
Eolymp #11429. Subordinates

Given the structure of a company, your task is to calculate for each employee the number of their subordinates.

Input. The first line has an integer n (1≤n≤2⋅105) — the number of employees. The employees are numbered 1, 2, ..., n, and employee 1 is the general director of the company.

After this, there are n−1 integers: for each employee 2, 3, ... , n their direct boss in the company.

Output. Print n integers: for each employee 1, 2, ..., n the number of their subordinates.

Sample input
5
1 1 2 3
Sample output
4 1 1 0 0
Open problem
Solution

The structure of the company is a tree. For each vertex v in the tree, determine the number of vertices sum[v] in the subtree where v is the root. The number of subordinates for employee v will be sum[v]−1 (since employee v is not his own subordinate).

Initiate a depth-first search starting from vertex 1, which corresponds to the general director of the company. If to1​,to2​,...,tok​ are the children of vertex v, then

sum[v]=1+sum[to1​]+sum[to2​]+...+sum[tok​]

If vertex v is a leaf in the tree, set sum[v]=1.

Example

Consider the example. Next to each employee (vertex of the graph), the number of their subordinates is indicated.

Algorithm implementation

The tree is stored in the adjacency list g.

vector<vector<int> > g;

The dfs function performs a depth-first search from vertex v and returns the number of vertices in the subtree rooted at v. Vertex p is the parent of v in the depth-first search.

int dfs(int v, int p)
{

Initially, the subtree rooted at v contains only the vertex v.

  sum[v] = 1;

Consider the edge v→to. If to=p, add to sum[v] the number of vertices in the subtree rooted at to.

  for (int to : g[v])
    if (to != p) sum[v] += dfs(to, v);

Return the number of vertices in the subtree rooted at v.

  return sum[v];
}

The main part of the program. Read the input data. Construct a tree.

scanf("%d", &n);
g.resize(n + 1);

for (i = 2; i <= n; i++)
{
  scanf("%d", &x);

For employee i, its immediate supervisor in the company is x.

  g[i].push_back(x);
  g[x].push_back(i);
}

Start the depth-first search from vertex 1.

sum.resize(n + 1);
dfs(1, -1);

Print the answer. The number of subordinates of employee i is equal to sum[i]–1.

for (i = 1; i <= n; i++)
  printf("%d ", sum[i] - 1);
printf("\n");

Java implementation

import java.util.*;

public class Main
{
  static ArrayList<Integer>[] g;	
  static int sum[];
  static int n;
  
  static int dfs(int v, int p)
  {
    sum[v] = 1;
    for(int to : g[v])
      if (to != p) sum[v] += dfs(to,v);
    return sum[v];
  }
  
  @SuppressWarnings("unchecked")
  public static void main(String[] args)
  {
    Scanner con = new Scanner(System.in);
    n = con.nextInt();
    g = new ArrayList[n+1]; // unchecked
    sum = new int[n+1];

    for (int i = 0; i <= n; i++)
      g[i] = new ArrayList<Integer>();

    for (int i = 2; i <= n; i++)
    {
      int x = con.nextInt();
      g[i].add(x);
      g[x].add(i);
    }
    
    dfs(1,-1);
    for (int i = 1; i <= n; i++)
      System.out.print(sum[i] - 1 + " ");    
    System.out.println();
    con.close();
  }
}   
Java
42 lines
842 bytes

Python implementation

Increase the stack for recursion.

import sys
sys.setrecursionlimit(1000000)

The dfs function performs a depth-first search from vertex v and returns the number of vertices in the subtree rooted at v. Vertex p is the parent of v in the depth-first search.

def dfs(v, p):

Initially, the subtree rooted at v contains only the vertex v.

  sum[v] = 1

Consider the edge v→to. If to=p, add to sum[v] the number of vertices in the subtree rooted at to.

  for to in g[v]:
    if to != p: sum[v] += dfs(to,v)

Return the number of vertices in the subtree rooted at v.

  return sum[v]

The main part of the program. Read the number of employees.

n = int(input())

Initialize the adjacency list of the graph g.

g = [[] for _ in range(n+1)]

Read the input data. Construct a tree.

lst = list(map(int, input().split()))
for i in range(2,n+1):
  a = lst[i-2]

For employee i (i≥2), its immediate supervisor in the company is a.

  g[a].append(i)
  g[i].append(a)

Initialize the list sum.

sum = [0] * (n + 1)

Start the depth-first search from vertex 1.

dfs(1,-1)

Print the answer. The number of subordinates of employee i is equal to sum[i]−1.

for i in range(1, n+1):
  print(sum[i] - 1,end=" ")
Eolymp #11428. Tree Diameter

A tree consisting of n vertices is given.

The diameter of a tree is the maximum distance between two vertices. Find the diameter of the given tree.

Input. The first line contains an integer n (1≤n≤2⋅105) — the number of vertices in the tree. The vertices are numbered from 1 to n.

Each of the following n−1 lines describes an edge and contains two integers a and b (1≤a,b≤n), indicating that there is an edge between vertices a and b.

Output. Print one integer — the diameter of the tree.

Sample input
5
1 2
1 3
3 4
3 5
Sample output
3
Open problem
Solution

Perform a depth-first search starting from any vertex, for example, from vertex 1. Find the vertex farthest from 1 and denote it as vertex a. Then, start a depth-first search from vertex a to find the vertex that is farthest from a, which we'll denote as vertex b. The distance between vertices a and b will then be the maximum possible and equal to the diameter of the tree.

Example

Consider the tree from the example.

The diameter of the tree is 3. The maximum distance is achieved, for example, between vertices 2 and 5.

Exercise

Find the diameter of the tree.

Algorithm implementation

Store the input tree in the adjacency list g. Store the shortest distances from the source to each vertex in the array dist.

vector<vector<int>> g;
vector<int> dist;

The dfs function performs a depth-first search from vertex v. The shortest distance from the source to vertex v is d. The parent of vertex v in the depth-first search is p.

void dfs(int v, int d, int p = -1)
{

Store the shortest distance from the source to vertex v in dist[v].

  dist[v] = d;

Iterate over all edges (v,to). For each son to of vertex v (where to is not a parent of v) initiate a depth-first search. The distance from the source to vertex to will be d+1.

  for (int to : g[v])
    if (to != p) dfs(to, d + 1, v);
}

The main part of the program. Read the input data.

scanf("%d", &n);
g.resize(n + 1);
dist.resize(n + 1);

Construct an undirected tree.

for (i = 1; i < n; i++)
{
  scanf("%d %d", &u, &v);
  g[u].push_back(v);
  g[v].push_back(u);
}

Find the shortest distances from vertex 1 to all other vertices. The vertex farthest from it is a.

dfs(1, 0, -1);
a = max_element(dist.begin(), dist.begin() + n + 1) – 
                dist.begin();

Find the shortest distances from vertex b to all other vertices. The vertex farthest from it is b.

dfs(a, 0,  -1);
b = max_element(dist.begin(), dist.begin() + n + 1) – 
                dist.begin();

Print the diameter of the tree — the shortest distance from a to b.

printf("%d\n", dist[b]);
Eolymp #11430. Tree Distances I

A tree consisting of n vertices is given.

For each vertex, determine the maximum distance to any other vertex.

Input. The first line contains an integer n (1≤n≤2⋅105) — the number of vertices in the tree. The vertices are numbered from 1 to n.

The next n−1 lines describe the edges: each line contains two integers a and b (1≤a,b≤n), indicating that there is an edge between vertices a and b.

Output. Print n integers. For each vertex from 1 to n, print the maximum distance to any other vertex in the tree.

Sample input
5
1 2
1 3
3 4
3 5
Sample output
2 3 2 3 3
Open problem
Solution

Using two depth-first search (DFS) runs, we find the diameter of the tree.

Since a tree is a connected graph without cycles, both depth-first search (DFS) and breadth-first search (BFS) correctly find the shortest distances from the starting vertex to all others.

Let a and b be two vertices that are the farthest apart — that is, they define the diameter of the tree.

Compute the shortest distances from vertex a to all other vertices and store them in the array dista, and from vertex b — in the array distb.

Then, for each vertex i, the greatest distance to any other vertex is equal to:

max(dista[i],distb[i])

Example

Consider a tree from the example.

The diameter of the tree can be, for example, the path between vertices 2 and 5.

Algorithm implementation

The input tree is stored as an adjacency list g.

The shortest distances from vertices a and b are stored in the arrays dista and distb, respectively.

vector<vector<int>> g;
vector<int> dista, distb;

The function dfs performs a depth-first search starting from vertex v.

  • The parameter d represents the shortest distance from the source to vertex v.

  • The results are stored in the array dist.

  • The parameter p indicates the parent of vertex v during the traversal.

void dfs(int v, int d, vector<int>& dist, int p = -1)
{

Store the shortest distance from the source to vertex v in dist[v].

  dist[v] = d;

Iterate over all edges (v,to). For each child to of vertex v (i.e., a vertex to that is not the parent of v), perform a depth-first search.

  for (int to : g[v])
    if (to != p) dfs(to, d + 1, dist, v);
}

The main part of the program. Read the input data.

scanf("%d", &n);
g.resize(n + 1);
dista.resize(n + 1);
distb.resize(n + 1);

Construct an undirected tree.

for (i = 1; i < n; i++)
{
  scanf("%d %d", &u, &v);
  g[u].push_back(v);
  g[v].push_back(u);
}

Compute the shortest distances from vertex 1. The vertex farthest from it is vertex a.

dfs(1, 0, dista, -1);
a = max_element(dista.begin() + 1, dista.begin() + n + 1) - dista.begin();

Next, compute the shortest distances from vertex a and store them in the array dista. The vertex farthest from a turns out to be vertex b. The distance between vertices a and b is the diameter of the tree.

dfs(a, 0, dista, -1);
b = max_element(dista.begin() + 1, dista.begin() + n + 1) - dista.begin();

Then compute the shortest distances from vertex b and store them in the array distb.

dfs(b, 0, distb, -1);

For each vertex i, print the distance to the farthest vertex from it.

for (i = 1; i <= n; i++)
  printf("%d ", max(dista[i], distb[i]));
printf("\n");

Python implementation

Increase the recursion stack size.

import sys
sys.setrecursionlimit(1000000)

The function dfs performs a depth-first search starting from vertex v.

  • The parameter d represents the shortest distance from the source to vertex v.

  • The results are stored in the array dist.

  • The parameter p indicates the parent of vertex v during the traversal.

def dfs(v, d, dist, p =- 1):

Store the shortest distance from the source to vertex v in dist[v].

  dist[v] = d

Iterate over all edges (v,to). For each child to of vertex v (i.e., a vertex to that is not the parent of v), perform a depth-first search.

  for to in g[v]:
    if to != p:
      dfs(to, d + 1, dist, v)

The main part of the program. Read the input data.

n = int(input())
g = [[] for _ in range(n + 1)]
dista = [0] * (n + 1)
distb = [0] * (n + 1)

Construct an undirected tree.

for _ in range(n - 1):
  u, v = map(int, input().split())
  g[u].append(v)
  g[v].append(u)

Compute the shortest distances from vertex 1. The vertex farthest from it is vertex a.

dfs(1, 0, dista)
a = dista.index(max(dista[1:]))

Next, compute the shortest distances from vertex a and store them in the array dista. The vertex farthest from a turns out to be vertex b. The distance between vertices a and b is the diameter of the tree.

dfs(a, 0, dista)
b = dista.index(max(dista[1:]))

Then compute the shortest distances from vertex b and store them in the array distb.

dfs(b, 0, distb)

For each vertex i, print the distance to the farthest vertex from it.

result = [max(dista[i], distb[i]) for i in range(1, n + 1)]
print(*result)
Eolymp #2157. Sum

Roman's parents gave him an undirected, connected, weighted graph with n vertices and n−1 edges. Roman wants to find the total length of all paths between every pair of vertices in the graph. The length of a path is defined as the sum of the weights of the edges it includes. Since the path from vertex u to vertex v is the same as the path from v to u, Roman treats them as a single path and does not distinguish between them.

Input. The first line contains an integer n (2≤n≤105) — the number of vertices in the graph.

Each of the next n−1 lines describes an edge and contains three integers: the numbers of the two vertices connected by the edge (vertices are numbered from 1 to n), and the weight of the edge.

Output. Print the total length of all distinct paths between all pairs of vertices, modulo 109.

Sample input 1
3
1 2 1
1 3 3
Sample output 1
8
Sample input 2
6
1 2 5
1 3 1
2 4 2
2 5 4
2 6 3
Sample output 2
90
Open problem
Solution

Perform a depth-first search starting from an arbitrary vertex of the tree. Consider an edge (u,v) with weight w. Let the number of vertices in the subtree rooted at v be c. Then, there are c vertices on one side of the edge and n−c vertices on the other side. The edge (u,v) is included in c⋅(n−c) distinct paths between pairs of vertices. Therefore, its contribution to the total length of all paths is c⋅(n−c)⋅w. We just need to iterate over all edges using depth-first search and sum up their contributions to the overall path length.

Example

The graphs given in the examples have the following structure:

Consider the contribution of the edges to the total length of all paths in the first example.

The edge (1,2) with weight 1 belongs to two paths:

  • 1−2;

  • 2−1−3;

Its contribution to the total sum is 1⋅2⋅1=2.

The edge (1,3) with weight 3 belongs to two paths:

  • 1−3;

  • 2−1−3;

Its contribution to the total sum is 2⋅1⋅3=6.

The total length of all paths is 2+6=8.

In the second example, let us consider the contribution of the edge (1,2) with weight 5 to the total length of all paths.

The number of vertices in the subtree rooted at vertex 2 is c=4 (including vertex 2 itself). Then, on the other side of the edge (1,2) there are n−c=6−4=2 vertices. Therefore, the number of distinct paths passing through the edge (1,2) is c⋅(n−c)=4⋅2=8.

Indeed, these are all pairs of vertices where one lies in the subtree of vertex 2, and the other lies outside of it:

1−2,1−2−4,1−2−5,1−2−6,3−1−2,3−1−2−4,3−1−2−5,3−1−2−6

The contribution of the edge (1,2) with weight 5 to the total length of all paths is

c⋅(n−c)⋅w=4⋅2⋅5=40
Algorithm implementation

Declare the value of the modulo as MOD.

#define MOD 1000000000

The input weighted graph is stored in an adjacency list g.

vector<vector<pair<int,int> > > g;

The function dfs performs a depth-first search starting from vertex v and returns the number of vertices in the subtree rooted at v (including v itself). The number of such vertices is counted in the variable cnt. Initially, set cnt=1 since the subtree already includes vertex v.

int dfs(int v, int p = -1)
{
  int cnt = 1;
  for (auto edge : g[v])
  {
    int to = edge.first;
    int w = edge.second;
C++
7 lines
130 bytes

Consider the edge (v,to) with weight w. Let c be the number of vertices in the subtree rooted at vertex to. Then, one side of the edge contains c vertices, and the other side contains n−c vertices. The edge (v,to) is included in c⋅(n−c) different paths between pairs of vertices. Thus, the contribution of the edge (v,to) to the total length of all paths is

c⋅(n−c)⋅w
    if (to != p)
    {
      int c = dfs(to, v);
      res = (res + 1LL * c * (n - c) * w) % MOD;
      cnt += c;
    }
  }
  return cnt;
}
C++
9 lines
149 bytes

The main part of the program. Read the input weighted graph into the adjacency list g.

scanf("%d", &n);
g.resize(n + 1);
for (i = 1; i < n; i++)
{
  scanf("%d %d %d", &u, &v, &d);
  g[u].push_back({v, d});
  g[v].push_back({u, d});
}
C++
8 lines
155 bytes

Start a depth-first search from vertex 1. The vertices of the graph are numbered from 1 to n.

dfs(1);

Print the answer.

printf("%lld\n",res);

Java implementation

import java.util.*;
import java.io.*;

class Edge 
{
  int to;
  int weight;
  public Edge(int to, int weight) 
  {
    this.to = to;
    this.weight = weight;
  }
}

public class Main
{
  static int MOD = 1000000000;
  static ArrayList<Edge>[] g;	
  static int n, m;
  static long res;
  
  static int dfs(int v, int p)
  {
    int cnt = 1;
    for(Edge edge : g[v])
    {
      int to = edge.to;
      int w = edge.weight;
      if (to != p)
      {
        int c = dfs(to, v);
        res = (res + 1L * c * (n - c) * w) % MOD;
        cnt += c;
      }
    }
    return cnt;
  }
  
  public static void main(String[] args)
  {
    FastScanner con = new FastScanner(System.in);
    n = con.nextInt();
    g = new ArrayList[n+1]; // unchecked

    for (int i = 0; i <= n; i++)
      g[i] = new ArrayList<Edge>();

    for (int i = 1; i < n; i++)
    {
      int u = con.nextInt();
      int v = con.nextInt();
      int d = con.nextInt();
      g[u].add(new Edge(v,d));
      g[v].add(new Edge(u,d));
    }
    
    dfs(1,-1);
    System.out.println(res);    
  }
}   

class FastScanner 
{
  BufferedReader br;
  StringTokenizer st;

  public FastScanner(InputStream inputStream) 
  {
    br = new BufferedReader(new InputStreamReader(inputStream));
    st = new StringTokenizer("");
  }

  public String next() 
  {
    while (!st.hasMoreTokens()) 
    {
      try 
      {
        st = new StringTokenizer(br.readLine());
      } catch (Exception e) 
      {
        return null;
      }
    }
    return st.nextToken();
  }
 
  public int nextInt() 
  {
    return Integer.parseInt(next());
  }
}
Java
92 lines
1694 bytes

Python implementation

Declare the value of the modulo as MOD.

MOD = 1000000000

The function dfs performs a depth-first search starting from vertex v and returns the number of vertices in the subtree rooted at v (including v itself). The number of such vertices is counted in the variable cnt. Initially, set cnt=1 since the subtree already includes vertex v.

def dfs(v, p = -1):
  global res
  cnt = 1
  for to, w in g[v]:

Consider the edge (v,to) with weight w. Let c be the number of vertices in the subtree rooted at vertex to. Then, one side of the edge contains c vertices, and the other side contains n−c vertices. The edge (v,to) is included in c⋅(n−c) different paths between pairs of vertices. Thus, the contribution of the edge (v,to) to the total length of all paths is

c⋅(n−c)⋅w
    if to != p:
      c = dfs(to, v)
      res = (res + c * (n - c) * w) % MOD
      cnt += c
  return cnt

The main part of the program. Read the input data.

n = int(input())
g = [[] for _ in range(n + 1)]
res = 0

for _ in range(n - 1):
  u, v, d = map(int, input().split())
  g[u].append((v, d))
  g[v].append((u, d))
Python
8 lines
170 bytes

Start a depth-first search from vertex 1. The vertices of the graph are numbered from 1 to n.

dfs(1)

Print the answer.

print(res)
Eolymp #2923. Tree

A rooted tree consisting of n vertices is given. Each vertex is painted in one of n colors. For each vertex v, determine the number of distinct colors that appear in the subtree rooted at v.

Input. The first line contains a single integer n (1≤n≤106). The following n lines describe the vertices of the tree. The i-th line contains two integers pi​ and ci​, where pi​ is the parent of vertex i, and ci​ is the color of vertex i (1≤ci​≤n). For the root of the tree, pi​=0.

Output. Print n integers — one for each vertex from 1 to n. For each vertex, print the number of distinct colors in the subtree rooted at that vertex.

Sample input
5
2 1
3 2
0 3
3 3
2 1
Sample output
1 2 3 1 1
Open problem
Solution

Perform a depth-first traversal of the tree starting from the root. For each vertex i, maintain a set si​, which accumulates the colors of all vertices in the subtree rooted at i.

If j is a child of vertex i during the traversal, then the set sj​ should be merged into si​.

The number of distinct colors in the subtree rooted at vertex i is equal to the size (cardinality) of the set si​.

Example

The color of each vertex is shown on the left. The set of colors in its subtree is shown on the right.

Algorithm implementation

The array color[i] stores the color of vertex i. The set s[i] will accumulate the colors that appear in the subtree rooted at vertex i.

The directed tree is represented as an adjacency list g. The array res[i] stores the number of distinct colors in the subtree rooted at vertex i.

#define MAX 1000010
int color[MAX], res[MAX];
set<int> s[MAX];
vector<vector<int> > g;

The function dfs performs a depth-first traversal of the tree starting from vertex v.

void dfs(int v)
{

First, the color of vertex v itself is added to the set s[v].

  s[v].insert(color[v]);

Iterate over the tree edges (v,to).

  for (int to : g[v])
  {
    dfs(to);

For each edge (v,to), merge the set s[to] into s[v]. If the size of s[v] is smaller than the size of s[to], swap them first. Then, the content of the smaller set s[to] is added to the larger set s[v].

    if (s[v].size() < s[to].size()) s[v].swap(s[to]);
    s[v].insert(s[to].begin(), s[to].end());

Clear the set s[to] — it is no longer needed.

    s[to].clear();
  }

After that, store the number of distinct colors in the subtree in res[v], which is the size of the set s[v].

  res[v] = s[v].size();
}

The main part of the program. Read the input data.

scanf("%d",&n);
g.resize(n+1);
for(i = 1; i <= n; i++)
{
  scanf("%d %d",&p,&c);
  g[p].push_back(i);
  color[i] = c;
}
C++
8 lines
128 bytes

Start the depth-first search from the root of the tree — the vertex with index 0.

dfs(0);

Print the answer.

for(i = 1; i <= n; i++)
  printf("%d ",res[i]);
printf("\n");
Eolymp #3983. Mancunian And Colored Tree

After a long and busy work week, the residents of Manchester and Liverpool decided to go hiking for the weekend. While walking through the forest, they came across a unique tree consisting of n vertices. The vertices of the tree are numbered from 1 to n, and each one is painted in one of c possible colors.

To fight off boredom, they decided to test their logical thinking. The root of the tree is vertex 1. For each vertex, the participants decided to find its nearest ancestor that shares the same color.

Input. The first line contains two integers n and c (1≤n,c≤105) — the number of vertices in the tree and the number of possible colors.

The second line contains n−1 integers: the i-th of them indicates the parent of vertex i+1.

The third line contains n integers — the colors of the vertices. Each color is an integer from 1 to c, inclusive.

Output. Print n integers in a single line: for each vertex, print the number of its nearest ancestor with the same color. If no such ancestor exists, print −1.

Sample input
5 4
1 1 3 3
1 4 2 1 2
Sample output
-1 -1 -1 1 3
Open problem
Solution

Let us consider a simplified version of the problem. Suppose all the vertices of the tree are painted the same color. We initialize an empty stack and perform a depth-first search starting from the root of the tree. When entering a vertex v, push it onto the stack. When the processing of vertex v is complete, remove the top element from the stack (which will be vertex v itself). After the DFS is finished, the stack will once again be empty.

Question: What will be on the top of the stack at the moment we enter vertex v, just before we push v onto the stack?

Now let us move on to solving the original problem. We create c stacks — one for each color (for example, using a vector of stacks). Initially, all stacks are empty. Perform a DFS starting from the root of the tree — vertex 1. The processing of a vertex v, painted in color color, includes the following steps:

  • If the stack s[color] is not empty, its top element contains the number of the vertex that is the nearest ancestor of vertex v with the same color. If the stack is empty, such a vertex does not exist, and the answer for vertex v should be −1.

  • Push vertex v onto the stack s[color].

  • Recursively perform a depth-first search from all children of vertex v.

  • After finishing the processing of vertex v, remove it from the stack s[color].

During the depth-first search, when we are at vertex v, the stacks contain information about the colors of all vertices along the unique path from the root to vertex v. In particular, the stack s[color] stores the numbers of all vertices on this path that are colored with color. These vertices are pushed onto the stack in the order they are visited during the DFS.

Example

When the depth-first search reaches vertex 5, two out of the c=4 stacks will be empty (those corresponding to colors 3 and 4).

Stack number 1 (for color 1) will contain vertex 1, and stack number 2 (for color 2) will contain vertex 3. Since vertex 5 has color 2, its nearest ancestor with the same color is at the top of stack number 2.

Therefore, the nearest ancestor of vertex 5 with the same color is vertex 3.

Algorithm implementation

Declare the arrays.

vector<int> col, res;
vector<vector<int> > g;
vector<stack<int> > s;

The dfs function performs a depth-first search starting from vertex v.

void dfs(int v)
{

Vertex v has color color.

  int color = col[v];

The number of the nearest ancestor of vertex v with the same color is stored in res[v].

  if(s[color].empty())
    res[v] = -1;
  else
    res[v] = s[color].top();

Push vertex v onto the stack corresponding to its color.

  s[color].push(v);

Iterate over all children to of vertex v and recursively perform a depth-first search from each of them.

  for (int to : g[v])
    dfs(to);

After processing vertex v is complete (i.e., upon exiting it), remove it from the stack corresponding to its color.

  s[color].pop();
}

The main part of the program. Read the input data. Construct a tree.

scanf("%d %d", &n, &c);
g.resize(n + 1);
for(i = 2; i <= n; i++)
{
  scanf("%d",&val);
  g[val].push_back(i);
}
C++
7 lines
119 bytes

Read the colors of the tree's vertices.

col.resize(n + 1);
for(i = 1; i <= n; i++)
  scanf("%d",&col[i]);

Start a depth first search from vertex 1.

s.resize(c + 1);
res.resize(n + 1);
dfs(1);

Print the answer.

for(i = 1; i <= n; i++)
  printf("%d ",res[i]);
printf("\n");

Python implementation

Increase the recursion stack size.

import sys
sys.setrecursionlimit(1000000)

The dfs function performs a depth-first search starting from vertex v.

def dfs(v):

Vertex v has color color.

  color = col[v]

The number of the nearest ancestor of vertex v with the same color is stored in res[v].

  if not s[color]:
    res[v] = -1
  else:
    res[v] = s[color][-1]

Push vertex v onto the stack corresponding to its color.

  s[color].append(v)

Iterate over all children to of vertex v and recursively perform a depth-first search from each of them.

  for to in g[v]:
    dfs(to)

After processing vertex v is complete (i.e., upon exiting it), remove it from the stack corresponding to its color.

  s[color].pop()

The main part of the program. Read the input data. Construct a tree.

n, c = map(int, input().split())
lst = list(map(int, input().split()))
g = [[] for _ in range(n + 1)]
for i in range(2, n + 1):
  val = lst[i - 2]
  g[val].append(i)

Read the colors of the tree's vertices.

col = [0] * (n + 1)
col[1:] = list(map(int, input().split()))

Start a depth first search from vertex 1.

s = [[] for _ in range(c + 1)]
res = [0] * (n + 1)
dfs(1)

Print the answer.

print(*res[1:])
Eolymp #10422. MooTube (Silver)

In his spare time, Farmer John created his own video-sharing platform called MooTube. On MooTube, his cows can record, publish, and discover lots of funny videos. So far, the cows have uploaded n videos, numbered from 1 to n.

However, Farmer John doesn't fully understand how to help his cows find new videos they might enjoy. He wants to implement a recommendation system for each video — a list of "recommended videos" based on their similarity to videos already watched.

To determine how similar two videos are, John introduces a relevance metric. He manually evaluates n−1 pairs of videos and assigns each pair a relevance score. Using these pairs, John builds a network where each video is a node in a graph, and the selected pairs are connected by edges with the given relevance values.

For convenience, John selects the n−1 pairs in such a way that there is exactly one path between any two videos — in other words, the video network forms a tree.

He decides that the relevance between two videos should be defined as the minimum relevance among all edges along the path connecting them.

Now, Farmer John wants to choose a value k such that, for any video on MooTube, all other videos with relevance at least k to it will be displayed as recommendations. However, he worries that too many recommendations might distract his cows from producing milk, so he wants to precisely estimate how many videos would be recommended for various values of k.

Farmer John turns to you for help: you are given several queries, each consisting of a value k and a video number. For each query, you need to determine how many other videos would be recommended if the minimum required relevance is k.

Input. The first line contains two integers n and q (1≤n,q≤5000) — the number of videos and the number of queries, respectively.

The next n−1 lines describe pairs of videos whose relevance has been manually evaluated by Farmer John. Each of these lines contains three integers pi​,qi​ и ri​ (1≤pi​,qi​≤n,1≤ri​≤109), , meaning that videos pi​ and qi​ are connected by an edge with relevance ri​.

Then follow q lines, each describing one of Farmer John’s queries. The i-th line contains two integers ki​ and vi​ (1≤ki​≤109,1≤vi​≤n) — meaning that in the i-th query, Farmer John wants to know how many videos will be recommended for video vi​ if the minimum acceptable relevance for recommendations is k=ki​.

Output. Print q lines. In the i-th line, output the answer to Farmer John's i-th query.

Example. Farmer John has established the following connections between videos:

  • Video 1 and video 2 have a relevance of 3,

  • Video 2 and video 3 have a relevance of 2,

  • Video 2 and video 4 have a relevance of 4.

Based on this data, we can compute the relevance between other pairs of videos:

  • Video 1 and video 3: relevance = min(3,2)=2,

  • Video 1 and video 4: relevance = min(3,4)=3,

  • Video 3 and video 4: relevance = min(2,4)=2.

Now let’s see which videos will be recommended for the following queries:

  • From video 2 with k=1, videos 1,3, and 4 will be recommended.

  • From video 1 with k=3, videos 2 and 4 will be recommended.

  • From video 1 with k=4, no videos will be recommended.

Sample input
4 3
1 2 3
2 3 2
2 4 4
1 2
4 1
3 1
7 lines
41 bytes
Sample output
3
0
2
Open problem
Solution

For each query given by the pair (ki​,vi​), perform a depth-first or breadth-first search starting from vertex vi​. During the traversal, only follow edges whose relevance is at least ki​.

Count the number of reachable vertices res, excluding the starting vertex vi​. The value of res is the answer to the query (ki​,vi​).

Example

Farmer John has established the following connections between videos:

  • Video 1 and video 2 have a relevance of 3,

  • Video 2 and video 3 have a relevance of 2,

  • Video 2 and video 4 have a relevance of 4.

Based on this data, we can compute the relevance between other pairs of videos:

  • Video 1 and video 3: relevance = min(3,2)=2,

  • Video 1 and video 4: relevance = min(3,4)=3,

  • Video 3 and video 4: relevance = min(2,4)=2.

Now let’s see which videos will be recommended for the following queries:

  • From video 2 with k=1, videos 1,3, and 4 will be recommended.

  • From video 1 with k=3, videos 2 and 4 will be recommended.

  • From video 1 with k=4, no videos will be recommended.

Algorithm implementation

Store the input graph in an adjacency list g.

vector<vector<pair<int, int>>> g;

The bfs function performs a breadth-first search starting from the vertex start and returns the number of visited vertices, excluding the vertex start itself. During the search, we only move along edges whose weight is at least the given value k.

int bfs(int start, int k)
{

The shortest distance array is not needed in this problem — it's sufficient to simply track whether a vertex is visited. If vertex i is visited, set used[i]=1.

  vector<int> used(n + 1);
  used[start] = 1;

Initialize a queue and add the starting vertex start to it.

  queue<int> q;
  q.push(start);

The number of visited vertices during the breadth-first search (excluding the starting vertex start) is counted in the variable res.

  int res = 0;

Start the breadth-first search.

  while (!q.empty())
  {

Extract the vertex v from the queue q.

    int v = q.front(); q.pop();

Iterate over all edges adjacent to vertex v.

    for (auto edge : g[v])
    {

Consider the edge (v,to) with weight kto.

     int to = edge.first;
     int kto = edge.second;

If the edge weight kto is less than k, skip this edge (moving along it is forbidden).

      if (kto < k) continue;

If the vertex to is not visited yet, add it to the queue, mark it as visited, and increment the counter res by 1.

      if (used[to] == 0)
      {
        q.push(to);
        used[to] = 1;
        res++;
      }
    }
  }
C++
8 lines
116 bytes

Return the answer to the query.

  return res;
}

The main part of the program. Read the input data.

scanf("%d %d", &n, &qu);
g.resize(n + 1);
for (i = 0; i < n - 1; i++)
{
  scanf("%d %d %d", &p, &q, &r);

Add an undirected edge (p,q) with weight r to the adjacency list.

  g[p].push_back({ q, r });
  g[q].push_back({ p, r });
}

Process the qu queries one by one.

while (qu--)
{
  scanf("%d %d", &k, &v);

Perform a breadth-first search (BFS) starting from vertex v and print the number of vertices visited during the BFS.

  printf("%d\n", bfs(v, k));
}

Algorithm implementation — depth first search

Store the input graph in an adjacency list g.

vector<vector<pair<int, int>>> g;

The dfs function performs a depth-first search starting from vertex v and returns the number of visited vertices. During the search, we only traverse edges whose weight is not less than the given value k.

int dfs(int v, int k)
{

Mark vertex v as visited.

  used[v] = 1;

Count the number of vertices visited in the subtree rooted at v during the depth-first search in the variable res.

  int res = 1;

Iterate over all edges adjacent to vertex v.

  for (auto edge : g[v])
  {

Consider the edge (v,to) with weight kto.

    int to = edge.first;
    int kto = edge.second;

If the weight kto is less than k, skip this edge (moving along it is forbidden).

    if (kto < k) continue;

If vertex to is not visited yet, perform a depth-first search starting from it. The number of vertices visited in the subtree rooted at to is added to res.

    if (used[to] == 0)
      res += dfs(to, k);
  }
  return res;
}

The main part of the program. Read the input data.

scanf("%d %d", &n, &qu);
g.resize(n + 1);
for (i = 0; i < n - 1; i++)
{
  scanf("%d %d %d", &p, &q, &r);

Add an undirected edge (p,q) with weight r to the adjacency list.

  g[p].push_back({ q, r });
  g[q].push_back({ p, r });
}

Process the qu queries one by one.

while (qu--)
{
  scanf("%d %d", &k, &v);

Perform a depth-first search starting from vertex v. Print the number of vertices visited during the depth-first search, excluding the starting vertex v.

  used.assign(n + 1, 0);
  printf("%d\n", dfs(v, k) - 1);
}
Eolymp #10653. XOR Sum

The tree with n vertices is given. The edges of the tree have weights of only 0 or 1. Let’s find the XOR sum between all pairs of vertices. Compute the sum of all XOR sums.

Input. The first line contains the number of vertices n (2≤n≤105) in the graph. The next n−1 lines describe the edges. Each line contains three integers: the numbers of the vertices connected by the edge (vertices are numbered from 1 to n) and the weight of the edge (0 or 1).

Output. Print the sum of XOR sums between all pairs of vertices.

Sample input
5
1 2 1
2 3 1
2 4 0
4 5 1
Sample output
6
Open problem
Solution

Using depth-first search, compute the XOR sum between vertex 1 and all other vertices. Store the XOR sum between vertex 1 and v in x[v]. Let ones be the number of ones and zeroes be the number of zeros in the array x (ones+zeroes=n). Then the answer to the problem will be ones⋅zeroes.

Example

Consider the graph provided in the example.

Next to each vertex v, the XOR sum x[v] between 1 and v is written. If x[v]=x[u] for some vertices v and u, then the XOR sum between them is equal to 0, thus contributing 0 to the total sum. The XOR sum for each pair of vertices (v,u), for which x[v]=x[u], contributes 1 to the total sum.

Therefore, the answer is equal to the number of pairs of vertices (v,u) for which x[v]=x[u]. This number is equal to ones⋅zeroes=2⋅3=6. The pairs of vertices that contribute 1 to the total sum are: (1,2),(1,4),(2,3),(2,5),(3,4),(4,5).

Algorithm implementation

Store the input graph in the adjacency list g. Declare an array x.

vector<vector<pair<int,int> > > g;
vector<int> x;

The dfs function implements a depth-first search that computes the XOR sum x[v] between vertices 1 and v. The current XOR sum between 1 and v is represented by cur_xor. The parent of vertex v is p.

void dfs(int v, int cur_xor, int p = -1)
{
  x[v] = cur_xor;
  for (auto z : g[v])
  {
    int to = z.first;
    int w = z.second;
    if (to != p) dfs(to, cur_xor ^ w, v);
  }
}
C++
10 lines
189 bytes

The main part of the program. Read the input graph.

scanf("%d", &n);
g.resize(n + 1);
x.resize(n + 1);

for (i = 1; i < n; i++)
{
  scanf("%d %d %d", &u, &v, &d);
  g[u].push_back({v, d});
  g[v].push_back({u, d});
}
C++
10 lines
175 bytes

Start the depth-first search from vertex 1.

dfs(1, 0, -1);

Compute the number of zeroes and ones in the array x.

ones = 0;
for (i = 1; i <= n; i++)
  if (x[i] == 1) ones++;
zeroes = n - ones;

Print the answer.

printf("%lld\n", 1LL * ones * zeroes);

Python implementation

The dfs function implements a depth-first search that computes the XOR sum x[v] between vertices 1 and v. The current XOR sum between 1 and v is represented by cur_xor. The parent of vertex v is p.

def dfs(v, cur_xor = 0, p = -1):
  x[v] = cur_xor
  for to, w in g[v]:
    if to != p:
      dfs(to, cur_xor ^ w, v)

The main part of the program. Read the input graph.

n = int(input())
g = [[] for _ in range(n + 1)]
for _ in range(n - 1):
  u, v, d = map(int, input().split())
  g[u].append((v, d))
  g[v].append((u, d))

Initialize the list x.

x = [0] * (n + 1)

Start the depth-first search from vertex 1.

dfs(1)

Compute the number of zeroes and ones in the array x.

ones = sum(x)
zeroes = n - ones

Print the answer.

print(ones * zeroes)
Eolymp #10654. Unique Color

A tree with n vertices numbered from 1 to n is given. The i-th edge of the tree connects vertices ai​ and bi​. Vertex i is colored with color ci​ (in this problem, colors are represented by integers).

A vertex x is considered good if the shortest path from vertex 1 to vertex x does not contain any other vertex with the same color as vertex x, except for x itself.

Find all good vertices.

Input. The first line contains an integer n (2≤n≤105) — the number of vertices.

The second line contains n integers — the colors of the vertices: c1​,c2​,…,cn​ (1≤ci​≤105).

Each of the following n−1 lines contains two integers ai​ and bi​ (1≤ai​,bi​≤n) — the edges of the tree.

Output. Print all good vertices, one per line, in ascending order of their indices.

Sample input 1
6
2 7 1 8 2 8
1 2
3 6
3 2
4 3
2 5
7 lines
41 bytes
Sample output 1
1
2
3
4
6
Sample input 2
10
3 1 4 1 5 9 2 6 5 3
1 2
2 3
3 4
4 5
5 6
6 7
7 8
8 9
9 10
11 lines
71 bytes
Sample output 2
1
2
3
5
6
7
8
7 lines
21 bytes
Open problem
Solution

Start a depth-first search from vertex 1. When entering vertex v with color color[v], increment the value of used[color[v]] by 1. The value used[color[v]] represents how many times the color color[v] has appeared on the path from vertex 1 to vertex v (inclusive). If this color appears exactly once on the path, then vertex v is considered good, and we add it to the result set.

When exiting vertex v, decrement the value of used[color[v]] by 1.

Example

The graph from the first example looks as follows:

Vertex 5 is not good because on the path 1−2−5, vertices 5 and 1 have the same color.

Vertex 6 is good because on the path 1−2−3−6, the colors of all intermediate vertices differ from the color of vertex 6.

Algorithm implementation

Store the input graph in the adjacency list g. Declare the arrays.

vector<int> used, color;
vector<vector<int>> g;
set<int> st;

The dfs function implements depth-first search. The variable par is the parent of vertex v.

void dfs(int v, int par)
{

Vertex v has the color color[v]. Mark that a vertex with this color is visited on the path from vertex 1.

  used[color[v]]++;

The value used[color[v]] indicates how many times the color color[v] has appeared on the path from vertex 1 to vertex v (inclusive). If this color appears exactly once on the path, then vertex v is considered good, and we add it to the result set st.

  if (used[color[v]] == 1) st.insert(v);

Iterate over all edges (v,to) originating from vertex v. If vertex to is not the parent of v (to=par), start a depth-first search from to. In this case, the parent of to becomes v.

  for (int to : g[v])
    if (to != par) dfs(to, v);

When exiting vertex v, decrement the value of used[color[v]] by 1.

  used[color[v]]--;
}

The main part of the program. Read the number of vertices n and the array of colors.

scanf("%d", &n);
color.resize(n + 1);
for (i = 1; i <= n; i++)
  scanf("%d", &color[i]);

Read the graph.

used.resize(100001);
g.resize(n + 1);
for (i = 1; i < n; i++)
{
  scanf("%d %d", &x, &y);
  g[x].push_back(y);
  g[y].push_back(x);
}
C++
8 lines
142 bytes

Start a depth-first search from vertex 1.

dfs(1, 1);

Print the good vertices. They are contained in the set st.

for (int val : st)
  printf("%d\n", val);
Eolymp #10655. Virus Tree 2

You are given a tree with n vertices and n−1 edges. The vertices are numbered from 1 to n, and the i-th edge connects vertices ai​ and bi​.

You have k colors available. For each vertex in the tree, you choose one of the k colors to paint it, subject to the following condition:

  • If the distance between two different vertices x and y is less than or equal to two, then x and y have different colors.

How many ways are there to color the tree? Find the answer modulo 109+7.

Input. The first line contains two numbers n and k (1≤n,k≤105). Each of the following n−1 lines contains two integers ai​ and bi​ (1≤ai​,bi​≤n).

Output. Print the number of ways to color the tree modulo 109+7

Sample input 1
4 3
1 2
2 3
3 4
Sample output 1
6
Sample input 2
5 4
1 2
1 3
1 4
4 5
Sample output 2
48
Open problem
Solution

Let's start coloring the tree from vertex 1. It can be colored with k colors.

Suppose we are at vertex v. If it has no parent (v=1), then its children can be colored with k−1 colors. If v has a parent, then its children can be colored with k−2 colors. Since the distance between the children of vertex v is 2, they cannot be colored with the same color.

If vertex v has a parent, then the first child of v can be colored with k−2 colors, the second child with k−3 colors, and so on.

It remains to multiply the numbers of colors that can be used to color each vertex.

Example

For the first example, there are 6 ways of coloring:

Let's consider the second test. Next to each vertex, we'll indicate the number of colors it can be colored with. For example, vertex 5 can be colored with 2 colors, as its color must not match the colors of vertices 1 and 4. The total number of ways to color the tree is equal to 4⋅3⋅2⋅1⋅2=48.

Algorithm implementation

Store the input graph in the adjacency list g. Declare a constant MOD for the modulo.

#define MOD 1000000007
vector<vector<int>> g;

The dfs function returns the number of ways to color the tree rooted at vertex v modulo MOD=109+7. The parent of vertex v is p. Vertex v can be colored with col colors.

int dfs(int v, int col, int p = -1)
{
  long long res = col;

We are at vertex v. In the variable cnt, we keep track of the number of colors that cannot be used to paint the children of vertex v. Initially, set cnt=1, since the children of vertex v cannot have the same color as vertex v. If vertex v has a parent (p=−1), then the children of vertex v also cannot have the color of v's parent.

  int cnt = 1;
  if (p >= 0) cnt++;

Iterate over the sons to of the vertex v.

  for (int to : g[v])
    if (to != p)
    {

The vertex to can be colored with k−cnt colors. With each son, the cnt value will be increased by 1.

      res = (res * dfs(to, k - cnt, v)) % MOD;
      cnt++;
    }
  return res;
}

The main part of the program. Read the input data.

scanf("%d %d", &n, &k);
g.resize(n + 1);
for (i = 0; i < n - 1; i++)
{
  scanf("%d %d", &a, &b);
  g[a].push_back(b);
  g[b].push_back(a);
}
C++
8 lines
149 bytes

Start the depth-first search from vertex 1. Vertex number 1 can be colored with k colors.

res = dfs(1, k);

Print the answer.

printf("%lld\n", res);

Python implementation

Set the maximum depth of the Python interpreter stack to the specified limit.

import sys
sys.setrecursionlimit(1000000)

Declare a constant MOD for the modulo.

MOD = 1000000007

The dfs function returns the number of ways to color the tree rooted at vertex v modulo MOD=109+7. The parent of vertex v is p. Vertex v can be colored with col colors.

def dfs(v, col, p = -1):
  res = col

We are at vertex v. In the variable cnt, we keep track of the number of colors that cannot be used to paint the children of vertex v. Initially, set cnt=1, since the children of vertex v cannot have the same color as vertex v. If vertex v has a parent (p=−1), then the children of vertex v also cannot have the color of v's parent.

  cnt = 1
  if p >= 0: cnt += 1

Iterate over the sons to of the vertex v.

  for to in g[v]:
    if to != p:

The vertex to can be colored with k−cnt colors. With each son, the cnt value will be increased by 1.

      res = (res * dfs(to, k - cnt, v)) % MOD
      cnt += 1
  return res

The main part of the program. Read the input data.

n, k = map(int, input().split())
g = [[] for _ in range(n + 1)]
for _ in range(n - 1):
  a, b = map(int, input().split())
  g[a].append(b)
  g[b].append(a)

Start the depth-first search from vertex 1. Vertex number 1 can be colored with k colors.

res = dfs(1, k)

Print the answer.

print(res)
Eolymp #10684. Roads Improvement

There are n cities in the country and n−1 bidirectional roads, such that it is possible to travel from any city to any other city using only these roads. The cities are numbered with integers from 1 to n inclusive.

Initially, all roads are considered to be in bad condition, but the government plans to improve the state of some of them. The citizens will be satisfied with the improvements if there is no more than one bad road on the path from the capital, located in city 1, to any other city.

Determine the number of ways to improve the quality of some roads to meet this requirement. Since the answer can be very large, output it modulo 109+7.

Input. The first line contains a single integer n (2≤n≤2⋅105) — the number of cities in the country. The second line contains n−1 positive integers p2​,p3​,p4​,⋯,pn​ (1≤pi​≤i−1), describing the roads in the country. The number pi​ indicates that there is a road connecting city pi​ and city i.

Output. Print the number of ways to improve the quality of the roads modulo 109+7.

Sample input 1
3
1 1
Sample output 1
4
Sample input 2
6
1 2 2 1 5
Sample output 2
15
Open problem
Solution

Let f(v) denote the number of ways to improve the quality of roads for the subtree rooted at vertex v. Let to represent a child of vertex v.

  • If the road (v,to) remains in bad condition, then all roads in the subtree rooted at vertex to must be improved.

  • If the road (v,to) is improved, then the number of ways to improve the roads for the subtree rooted at vertex to is equal to f(to).

Let to1​,to2​,...,tok​ be the children of v. Then

f(v)=(f(to1​)+1)⋅(f(to2​)+1)⋅...⋅(f(tok​)+1)

If v is a leaf (a vertex with no children), then let f(v)=1.

Example

Consider the examples provided in the problem statement. For each vertex v, we'll record the value of f(v).

In the first example, there are 4 ways to improve the quality of roads such that on the path from city 1 to any other city, there is no more than one bad road. Bad roads are represented by dashed lines, while improved roads are represented by solid lines.

Algorithm implementation

All calculations are performed modulo MOD=109+7.

#define MOD 1000000007

The dfs function computes the value of f(v) — the number of ways to improve the quality of roads for the subtree rooted at vertex v. The parent of v is vertex p.

long long dfs(int v, int p = -1)
{
  long long s = 1;

If to1​,to2​,...,tok​ are the children of v, then

f(v)=(f(to1​)+1)⋅(f(to2​)+1)⋅...⋅(f(tok​)+1)
  for (int to : g[v])
    if (to != p)
    {
      long long sub = dfs(to, v);
      s = (s * (sub + 1)) % MOD;
    }
  return s;
}
C++
8 lines
140 bytes

Read the input data. Construct a tree.

scanf("%d", &n);
g.resize(n + 1);
for (i = 2; i <= n; i++)
{
  scanf("%d", &p);
  g[i].push_back(p);
  g[p].push_back(i);
}
C++
8 lines
132 bytes

Compute and print the result.

res = dfs(1);
printf("%lld\n", res);
Eolymp #11058. Chasing a butterfly

On clear summer days, Nyusha enjoys catching butterflies in the fresh air. But today, she encountered a particularly cunning butterfly — it flew into a labyrinth and tried to hide from her inside.

The labyrinth consists of n rooms, numbered from 1 to n, some of which are connected by corridors. It is known that there is exactly one simple path between any two rooms, passing only through the corridors. In other words, the labyrinth forms a tree with n vertices and n−1 edges.

The entrance to the labyrinth is located in room number 1. A leaf is defined as any room connected to exactly one other room and not being the root (i.e., not room 1). Each leaf contains an exit from the labyrinth. The butterfly starts its flight from room 1, heading toward one of the exits. It moves at a constant speed, never turns around, and travels through one corridor per minute, moving to a neighboring room. All corridors are of equal length.

To catch the butterfly, Nyusha decided to enlist the help of some friends. Initially, each of them can take position in any room that contains an exit. As soon as the butterfly begins its journey from the entrance toward some exit, the friends may immediately start moving from their rooms toward the entrance. They move at the same speed as the butterfly. If any of them encounters the butterfly — whether in a room or in the middle of a corridor — it is considered caught. However, if the butterfly reaches an exit without encountering any of the friends, it successfully escapes to freedom.

Help Nyusha determine the minimum number of friends needed to guarantee catching the butterfly, regardless of which exit it chooses.

Input. The first line contains a single integer n (2≤n≤200000) — the number of rooms in the labyrinth.

The following n−1 lines describe the corridors connecting the rooms. Each line contains two integers u and v (1≤u,v≤n,u=v) — the indices of the rooms connected by a corridor.

It is guaranteed that the given corridor system forms a tree.

Output. Print a single integer — the minimum number of friends required to guarantee that the butterfly will be caught.

Sample input 1
3
1 2
1 3
Sample output 1
2
Sample input 2
4
1 2
2 3
2 4
Sample output 2
1
Open problem
Solution

Perform a depth-first search starting from vertex 1. For each vertex, compute two values:

  • d[v] — the distance from vertex 1 to vertex v. If p is the parent of v, then

    d[v]=d[p]+1
  • h[v] — the distance from vertex v to the nearest leaf in the subtree rooted at v. If to1​,to2​,...,tok​ are the children of vertex v, then

    h[v]=1+min(h[to1​],h[to2​],...,h[tok​])
  • If v is a leaf of the tree, set h[v]=0.

Next, perform a second depth-first search. During this traversal, compute the value res[v] — the minimum number of friends that need to be placed in some of the leaves of the subtree rooted at vertex v in order to guarantee catching the butterfly, assuming it chooses to fly into this subtree.

If h[v]≤d[v], then res[v]=1. In this case, it is enough to place one friend in the leaf with the minimal depth in the subtree rooted at vertex v: they will reach vertex v no later than the butterfly flying from the root.

Otherwise, if to1​,to2​,...,tok​ are the children of vertex v, then

res[v]=res[to1​]+res[to2​]+...+res[tok​]

If the butterfly is not caught at vertex v itself, we must be prepared to intercept it in any of the subtrees of v's children.

The final answer to the problem is the value res[1].

Example

In the first example (shown in the left diagram), two friends are required, and they should be placed at the two exits. The butterfly can fly from vertex 1 either to vertex 2 or to vertex 3. In either case, it will be intercepted by one of Nyusha's friends.

In the second example (shown in the right diagram), one friend is enough, who can be placed at either of the two exits — vertex 3 or 4. The butterfly will fly from vertex 1 to vertex 2 in one minute. During that same time, the friend will reach vertex 2 and catch the butterfly there.

Let us consider the next example.

Nyusha will need three friends to catch the butterfly.

Algorithm implementation

Declare a constant for infinity and arrays.

#define INF 2000000000
vector<vector<int> > g;
vector<int> d, h, res;

The function dfs (v,p,cnt) performs a depth-first search starting from vertex v and returns the value h[v]. Here, p is the parent of vertex v, and cnt is the distance from vertex 1 to v. During the traversal, the values d[v] and h[v] are computed for each vertex v.

int dfs(int v, int p = -1, int cnt = 0)
{

The value cnt, which equals the distance from vertex 1 to v, is stored in d[v].

  d[v] = cnt;
  int height = INF;

Iterate over all the children of the vertex v. Consider the edge v→to. If to equals the parent of v (to=p), skip it and move to the next child.

  for (int to : g[v])
    if (to != p) 

In the variable height compute the value

min(h[to1​],h[to2​],...,h[tok​]),

where to1​,to2​,...,tok​ are the children of vertex v.

height=min(height,dfs(to,v,cnt+1));

If height=INF, then vertex v is a leaf, and set h[v]=0. Otherwise rerturn

h[v]=1+min(h[to1​],h[to2​],...,h[tok​])
  return h[v] = (height == INF) ? 0 : 1 + height;
}

The function dfs2 (v,p) performs a depth-first search starting from vertex v and returns the value res[v]. Here, p is the parent of vertex v.

int dfs2(int v, int p = -1)
{
  int s = 0;

Let to1​,to2​,...,tok​ be the children of vertex v. In the variable s, compute the sum:

res[to1​]+res[to2​]+...+res[tok​]
  for (int to : g[v])
    if (to != p)
    {
      dfs2(to, v);
      s += res[to];
    }

If h[v]≤d[v], then one friend is sufficient, and res[v]=1. Otherwise

res[v]=res[to1​]+res[to2​]+...+res[tok​]=s
  return res[v] = (h[v] <= d[v]) ? 1 : s;
}

The main part of the program. Read the input data. Construct a graph.

scanf("%d", &n);
g.resize(n + 1);
for (i = 0; i < n - 1; i++)
{
  scanf("%d %d", &a, &b);
  g[a].push_back(b);
  g[b].push_back(a);
}
C++
8 lines
142 bytes

Initialize the arrays.

d.resize(n + 1);
h.resize(n + 1);
res.resize(n + 1);

Run two depth-first searches. The first traversal computes the values of d[v] and h[v] for each vertex v.

dfs(1);
dfs2(1);

Print the answer — the minimum number of friends needed to catch the butterfly.

printf("%lld\n", res[1]);
Eolymp #11153. Kefa and Park

Kefa decided to celebrate his first big salary by going to a restaurant.

He lives near an unusual park. The park is a rooted tree with n vertices, rooted at vertex 1. Kefa's house is located at vertex 1 as well. Unfortunately for our hero, the park is inhabited by cats, and Kefa has already identified the vertices where the cats are located.

The leaf vertices of the park contain restaurants. Kefa wants to choose a restaurant, but he is very afraid of cats. Therefore, he will not go to a restaurant if, on the path from his house to the restaurant, he encounters more than m consecutive vertices with cats.

Your task is to help Kefa count the number of restaurants he can safely visit.

Input. The first line contains two integers n and m (2≤n≤105,1≤m≤n) — the number of vertices in the tree and the maximum number of consecutive vertices with cats that Kefa can tolerate.

The second line contains n integers a1​,a2​,...,an​, where each ai​ is either 0 (there is no cat at vertex i) or 1 (there is a cat at vertex i).

The following n−1 lines contains the tree's edges in the format xi​ yi​ (1≤xi​,yi​≤n,xi​=yi​), where xi​ and yi​ are vertices connected by an edge.

It is guaranteed that this set of edges forms a tree.

Output. Print the number of leaf vertices that Kefa can reach if there are no more than m consecutive vertices with cats on the way from his house.

Sample input 1
7 1
1 1 0 0 0 0 1
1 2
1 3
2 4
2 5
3 6
3 7
8 lines
50 bytes
Sample output 1
2
Sample input 2
8 2
1 1 0 1 0 1 0 1
1 2
2 3
2 5
2 6
3 4
6 7
6 8
9 lines
57 bytes
Sample output 2
2
Open problem
Solution

Start a depth first search from vertex 1, where Kefa's house is located. One of the parameters of the dfs function will store the number of consecutive vertices with cats. If at vertex v this number exceeds m, we'll not continue the depth-first search from this vertex. We count the number of visited leaf vertices. A vertex is a leaf if its degree is 1 and it is not a root (vertex 1).

Example

Let’s consider the trees from the first and second examples.

In the first example, Kefa can pass through only one consecutive vertex with a cat. He will be able to reach the restaurants located at vertices 6 and 7.

In the second example, Kefa can pass through two consecutive vertices with cats. He will be able to reach the restaurants located at vertices 4 and 5.

Algorithm implementation

The location of the cats is stored in the array cats. The tree is stored in the adjacency list g.

vector<int> cats;
vector<vector<int> > g;

The dfs function performs a depth-first search from vertex v and returns the number of restaurants that can be reached from this vertex (under the condition that no more than m consecutive vertices with cats are encountered on the path to them). The parent of vertex v is the vertex prev. The number of consecutive vertices with cats encountered up to and including vertex v is cnt.

int dfs(int v, int prev, int cnt)
{

If more than m consecutive cats are already encountered, further depth-first search from vertex v is not continued.

  if (cnt > m) return 0;

If we are at a leaf, increase the number of accessible restaurants.

  if (g[v].size() == 1 && prev != -1) return 1;

In the variable res count the number of restaurants reachable from vertex v.

  int res = 0;

Iterate over all edges (v,to) originating from vertex v. If the vertex to is a parent of v, do not move to this vertex.

  for (int to : g[v])
    if (to != prev) 
    {

If there is no cat at the vertex v, set c=0. Otherwise, assign c=cnt+1. The variable c stores the number of consecutive vertices with cats up to and including vertex to.

      int c = (cats[to] == 0) ? 0 : cnt + 1;

Start a depth-first search from vertex to.

      res += dfs(to, v, c);
    }

Return the result.

  return res;
}

The main part of the program. Read the input data.

scanf("%d %d", &n, &m);
cats.resize(n + 1);
for (i = 1; i <= n; i++)
  scanf("%d", &cats[i]);

Construct a tree.

g.resize(n + 1);
for (i = 1; i < n; i++)
{
  scanf("%d %d", &a, &b);
  g[a].push_back(b);
  g[b].push_back(a);
}
C++
7 lines
120 bytes

Print the result. The dfs function, called from vertex 1, returns the answer to the problem. Since vertex 1 has no parent, we pass the value −1 as the second parameter. Starting from vertex 1, we have encountered cats[1] cats.

printf("%d\n", dfs(1, -1, cats[1]));

Python implementation

The dfs function performs a depth-first search from vertex v and returns the number of restaurants that can be reached from this vertex (under the condition that no more than m consecutive vertices with cats are encountered on the path to them). The parent of vertex v is the vertex prev. The number of consecutive vertices with cats encountered up to and including vertex v is cnt.

def dfs(v, prev, cnt):

If more than m consecutive cats are already encountered, further depth-first search from vertex v is not continued.

  if cnt > m:
    return 0

If we are at a leaf, increase the number of accessible restaurants.

  if len(g[v]) == 1 and prev != -1:
    return 1

In the variable res count the number of restaurants reachable from vertex v.

  res = 0

Iterate over all edges (v,to) originating from vertex v. If the vertex to is a parent of v, do not move to this vertex.

  for to in g[v]:
    if to != prev:

If there is no cat at the vertex v, set c=0. Otherwise, assign c=cnt+1. The variable c stores the number of consecutive vertices with cats up to and including vertex to.

      c = 0 if cats[to] == 0 else cnt + 1

Start a depth-first search from vertex to.

      res += dfs(to, v, c)

Return the result.

  return res

The main part of the program. Read the input data.

n, m = map(int, input().split())
cats = [0] + list(map(int, input().split()))

Construct a tree.

g = [[] for _ in range(n + 1)]
for _ in range(n - 1):
  a, b = map(int, input().split())
  g[a].append(b)
  g[b].append(a)

Print the result. The dfs function, called from vertex 1, returns the answer to the problem. Since vertex 1 has no parent, we pass the value −1 as the second parameter. Starting from vertex 1, we have encountered cats[1] cats.

print(dfs(1, -1, cats[1]))
Eolymp #11609. Finding pairs

A tree with n vertices is given. Vertex 1 is considered the root. There is exactly one simple path between any two vertices.

Let d(i,j) denote the number of edges on the path from vertex i to vertex j.

Find the number of vertex pairs (i,j) such that the following equality holds:

d(i,j)=d(i,1)−d(j,1)

Input. The first line contains a single integer n (1≤n≤105) — the number of vertices in the tree.

Each of the next n−1 lines contains two integers, describing the edges of the tree: the pairs of vertices connected by an edge.

Output. Print a single integer — the number of pairs (i,j) for which d(i,j)=d(i,1)−d(j,1) holds.

Sample input
5
1 2
2 3
2 4
4 5
Sample output
13
Open problem
Solution

The condition d(i,j)=d(i,1)−d(j,1) holds for every pair of vertices (i,j) where j is an ancestor of i in the depth-first search starting from vertex 1.

Let us consider an example:

  • d(4,1)=2,d(4,1)−d(1,1)=2−0=2;

  • d(6,3)=2,d(6,1)−d(3,1)=3−1=2;

  • d(2,2)=0,d(2,1)−d(2,1)=1−1=0;

  • d(6,1)=3,d(6,1)−d(1,1)=3−0=3;

Let us define the depth h[v] of a vertex v as the number of vertices on the path from the root (vertex 1) to v. In the diagram, the depth is written next to each vertex.

Now, fix a vertex i and ask the question: how many vertices j are there such that the condition d(i,j)=d(i,1)−d(j,1) is satisfied?

For example, for i=6, there are 4 such vertices: j=1,3,4,6. Note that h[6]=4. Thus, for a fixed vertex i, there are exactly h[i] suitable vertices j.

To solve the problem, we need to compute the depth h[v] for each vertex v in the tree, and then find the sum of all these depths across the tree.

Example

Let us consider the graph from the example. Next to each vertex, we write its depth.

The sum of the depths of all vertices is 1+2+3+3+4=13.

Algorithm implementation

Define the graph's adjacency list as g. The depths of the vertices are stored in the array h.

vector<vector<int> > g;
vector<int> h;

The dfs function computes the depth of each vertex v. The parent of vertex v is the vertex p.

int dfs(int v, int p)
{
  for (int to : g[v])

Consider the edge (v,to). If the vertex to is not the parent of vertex v, compute its depth and perform a depth first search starting from it.

    if (to != p)
    {
      h[to] = h[v] + 1;
      dfs(to, v);
    }
  return h[v];
}
C++
7 lines
95 bytes

The main part of the program. Read the number of vertices n in the tree.

scanf("%d", &n);
g.resize(n + 1);

Construct a tree.

for (i = 2; i <= n; i++)
{
  scanf("%d %d", &a, &b);
  g[a].push_back(b);
  g[b].push_back(a);
}

Set the depth of vertex 1 to be 1.

h.resize(n + 1);
h[1] = 1;

Start a depth first search from the vertex 1.

dfs(1, -1);

Compute the answer — the sum of the depths of all vertices.

res = 0;
for (i = 1; i <= n; i++)
  res += h[i];

Print the answer.

printf("%lld\n", res);
Eolymp #11724. Infected Tree

You are given a binary tree, which is an acyclic connected undirected graph containing n vertices and n−1 edges. Each vertex has a degree of no more than 3. The root is the vertex number 1, with its degree not exceeding 2.

Unfortunately, the root of the tree is infected. The following process is repeated n times:

  • Huseyn either selects a vertex that is not yet infected (and not yet removed) and removes it along with all its incident edges, or he takes no action.

  • After that, the infection spreads to every vertex connected by an edge to an already infected vertex (all previously infected vertices remain infected).

Help Huseyn determine the maximum number of vertices he can save from infection (note that removed vertices are not considered saved).

Input. The first line contains the number of vertices in the tree n (2≤n≤3⋅105).

Each of the following n−1 lines contains two integers ui​ and vi​ (1≤ui​,vi​≤n), representing an edge of the tree.

It is guaranteed that the graph is a binary tree with the root at vertex 1.

Output. Print a single integer — the number of saved vertices.

Sample input 1
4
1 2
2 3
2 4
Sample output 1
2
Sample input 2
7
1 2
1 5
2 3
2 4
5 6
5 7
7 lines
33 bytes
Sample output 2
2
Open problem
Solution

For each vertex v, compute the number of vertices sum[v] in the subtree where it is the root. If to1​,to2​,...,tok​ are the children of vertex v, then

sum[v]=1+sum[to1​]+sum[to2​]+...+sum[tok​]

If vertex v is a leaf of the tree, then sum[v]=1.

Let dp[v] be the number of saved vertices in the subtree with root v, assuming that currently vertex v (and only it in the subtree) is infected.

If vertex v has only one child to, then dp[v]=sum[to]−1 (the removed vertex to is not considered saved).

Let vertex v have two children to1​ and to2​ (each vertex has a degree of at most 3). Then we have two options:

  • Remove to1​ and recursively save the vertices in the subtree with root to2​. The number of saved vertices will be sum sum[to1​]−1+dp[to2​].

  • Remove to2​ and recursively save the vertices in the subtree with root to1​. The number of saved vertices will be sum[to2​]−1+dp[to1​].

Among the two options, choose the one that maximizes the number of saved vertices.

Consider the depth-first search process when examining the edge (v,to).

Let a be the second child of vertex v. We need to compute the value dp[to]+sum[a]−1. Let's try to find it using only the vertices v and to. We know that:

sum[v]=1+sum[to]+sum[a]

From which it follows that:

sum[a]=sum[v]−sum[to]−1

Thus,

dp[to]+sum[a]−1=dp[to]+sum[v]−sum[to]−2

Iterate over the children to of the vertex v and compute:

dp[v]=max(dp[to]+sum[v]−sum[to]−2)

Example

In the first example, Huseyn is able to save two vertices, 3 and 4, by choosing vertex number 2 as his first move.|

Consider the second example. Next to each vertex v, place the labels sum[v]/dp[v]. The vertices chosen by Huseyn are displayed in green.

For example,

dp[1]=max(sum[2]−1+dp[5],sum[5]−1+dp[2])=max(2,2)=2,sum[1]=1+sum[2]+sum[5]=1+3+3=7
Algorithm implementation

Declare the arrays.

#define MAX 300005
int sum[MAX], dp[MAX];
vector<vector<int>> g;

The function dfs performs a depth-first search starting from vertex v. The parent of vertex v is p.

void dfs(int v, int p = -1)
{

Initialize the variables sum[v] and dp[v].

  sum[v] = 1;
  dp[v] = 0;

Perform a depth-first search starting from the children of vertex v.

  for (int to : g[v])
    if (to != p)
    {
      dfs(to, v);
      sum[v] += sum[to];
    }

Compute the value of dp[v].

  for (int to : g[v])
    if (to != p)
    {
      dp[v] = max(dp[v], sum[to] - 1); // if 1 son
      dp[v] = max(dp[v], dp[to] + sum[v] - sum[to] - 2);
    }
}
C++
7 lines
168 bytes

The main part of the program. Read the input data.

scanf("%d", &n);
g.clear(); g.resize(n + 1);
memset(dp, 0, sizeof(dp));
memset(sum, 0, sizeof(sum));
for (i = 1; i < n; i++)
{
  scanf("%d %d", &u, &v);
  g[u].push_back(v);
  g[v].push_back(u);
}
C++
10 lines
207 bytes

Perform a depth-first search starting from vertex 1.

dfs(1);

Print the answer.

printf("%d\n", dp[1]);

List of problems

  • 977. Is it a Tree?

  • 978. Get a tree

  • 11263. Even tree

  • 11429. Subordinates

  • 11428. Tree diameter

  • 11430. Tree distances I

  • 2157. Sum

  • 2923. Tree

  • 3983. Mancunian and colored tree

  • 10422. MooTube (Silver)

  • 10653. XOR Sum

  • 10654. Unique color

  • 10655. Virus tree 2

  • 10684. Roads improvement

  • 11058. Chasing a butterfly

  • 11153. Kefa and park

  • 11609. Finding pairs

  • 11724. Infected Tree


11

Comments (4)

Loading

Just a moment, getting data from the server
user897042•18 days ago

what a community i asked for help and noone cares.

1
user897042•19 days ago

https://codeforces.com/contest/315/problem/B how to solve this problem i wrote two updates one for assinging and second for interval with lazy. but it doesnt seem to work because two seperate updates doesnt have link between them.

1
medv•18 days ago

user897042 https://site.ada.edu.az/~medv/acm/Docs%20e-olimp/Volume%209/853_English.htm

1
user897042•19 days ago

user897042 need help on these two too:https://codeforces.com/contest/1857/problem/Ghttps://codeforces.com/contest/1702/problem/G

1