Mmedv1 year 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)$ 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$ n vertices, there are always exactly $n - 1$ 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 ( $D F S$ DFS, $L C A$ 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$ n vertices and $n - 1$ n - 1 edges.
An acyclic graph with $n - 1$ 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++

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 \leq n \leq 100)$ n~(1 \le n \le 100) in the graph. Then, an adjacency matrix of size $n \times n$ n \times n is given, where $1$ 1 indicates the presence of an edge, and $0$ 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.

Open problem

Declare the adjacency matrix of the graph $g$ g and the array $u se d$ used.

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

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

void dfs(int v, int p)
{

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

  if (flag) return;

Mark the vertex $v$ v as visited.

  used[v] = 1;

The vertex $v$ v is visited, increase $c$ c by $1$ 1.

  c++;

The edge $(v, i)$ (v, i) will be a back edge and form a cycle if $i \neq = p$ i \ne p and vertex $i$ i is already visited $(u se d [i] = 1)$ (used[i] = 1). If a cycle is detected, set $f l a g = 1$ flag = 1. If no cycle is detected, continue the search from vertex $i$ 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$ c. Set $f l a g = 0$ flag = 0 if there is no cycle in the graph. If a cycle is detected, $f l a g$ flag becomes $1$ 1.

c = 0;
flag = 0;

All vertices are initially unvisited (initialize the array $u se d$ used with zeroes).

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

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

dfs(0,-1);

The graph is not a tree if there is a back edge $(f l a g = 1)$ (flag = 1) or if the graph is not connected $(c \neq = n)$ (c \ne 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$ v. The parent of vertex $v$ v is $p$ p. If a cycle is detected, the $f l a g$ flag variable is set to $1$ 1.

def dfs(v, p):

Declare the global variables used by the function.

  global flag, c

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

  if flag: return

Mark the vertex $v$ v as visited.

  used[v] = 1

The vertex $v$ v is visited, increase $c$ c by $1$ 1.

  c += 1

  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.

n = int(input())

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

c = 0
flag = 0

All vertices are initially unvisited (initialize the list $u se d$ used with zeroes).

used = [0] * n

Create an adjacency matrix $g$ 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$ 0. Since it is the root of the tree, it has no parent. Pass the value $- 1$ -1 as the second argument to the dfs function.

dfs(0, -1)

The graph is not a tree if there is a back edge $(f l a g = 1)$ (flag = 1) or if the graph is not connected $(c \neq = n)$ (c \ne n).

if flag or c != n:
  print("NO")
else:
  print("YES")

Declare the adjacency matrix of the graph $g$ g and the array $u se d$ used.

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

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

void dfs(int v)
{

Mark the vertex $v$ v as visited.

  used[v] = 1;

Vertex $v$ v is visited, increase $c$ c by $1$ 1.

  c++;

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

There is an edge $(v, i)$ (v, i), that is $g [v] [i] = 1$ g[v][i] = 1;
The vertex $i$ i is not visited $(u se d [i] = 0)$ (used[i] = 0)

If both conditions are true, we continue the depth-first search from vertex $i$ 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$ n.

scanf("%d", &n);

Count the number of edges in the graph in the variable $E d g es$ Edges.

Edges = 0;

All vertices are initially unvisited (initialize the array $u se d$ used with zeroes).

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

Read the adjacency matrix $g$ 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)$ (u, v) and $(v, u)$ (v, u) are considered the same. Divide the value of $E d g es$ Edges by $2$ 2.

Edges /= 2;

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

c = 0;

Start the depth-first search from the vertex $1$ 1.

dfs(1);

The graph is a tree if the number of edges in it equals $n - 1$ n - 1, and also if it is connected $(c = n)$ (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$ v.

def dfs(v):
  global c

Mark the vertex $v$ v as visited.

  used[v] = 1

The vertex $v$ v is visited, increase $c$ c by $1$ 1.

  c += 1

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

There is an edge $(v, i)$ (v, i), that is $g [v] [i] = 1$ g[v][i] = 1;
The vertex $i$ i is not visited $(u se d [i] = 0)$ (used[i] = 0)

If both conditions are true, we continue the depth-first search from vertex $i$ 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.

n = int(input())

Count the number of edges in the graph in the variable $E d g es$ Edges.

Edges = 0

All vertices are initially unvisited (initialize the array $u se d$ used with zeroes).

used = [0] * (n + 1)

Create an adjacency matrix $g$ g, initially filled with zeroes.

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

Read the adjacency matrix $g$ 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)$ (u, v) and $(v, u)$ (v, u) are considered the same. Divide the value of $E d g es$ Edges by $2$ 2.

Edges //= 2

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

c = 0

Start the depth-first search from the vertex $1$ 1.

dfs(1)

The graph is a tree if the number of edges in it equals $n - 1$ n - 1, and also if it is connected $(c = n)$ (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 \leq n \leq 100)$ n~(1 \le n \le 100) and the number of edges $m$ m in the graph. The following $m$ m pairs of integers each represent an edge. It is guaranteed that the graph is connected.

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

Open problem

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

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

The dfs function performs a depth-first search starting from vertex $v$ v. Each edge $(v, i)$ (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++

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++

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$ 4 nodes shown below, you can remove at most $1$ 1 edge to create an even forest.

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

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

Open problem

Declare the arrays.

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

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

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

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

  int s = 1;

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

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

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

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

Return the number of vertices in the subtree rooted at $v$ 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$ 1. In the variable $r es$ res, count the number of removed edges.

res = 0;
dfs(1);

Since there is no edge from $1$ 1 to its parent, print $r es - 1$ res - 1.

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

Python implementation

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

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

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

  s = 1

Iterate over the children $i$ i of vertex $v$ v. Add the values of $df s (i)$ dfs(i) to $s$ 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$ s is even, remove the edge between $v$ v and its parent.

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

Return the number of vertices in the subtree rooted at $v$ 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

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

res = 0
dfs(1)

Since there is no edge from $1$ 1 to its parent, print $r es - 1$ 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 \leq n \leq 2 \cdot 1 0^{5})$ n~(1 \le n \le 2 \cdot 10^5) — the number of employees. The employees are numbered $1$ 1, $2$ 2, ..., $n$ n, and employee $1$ 1 is the general director of the company.

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

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

Open problem

The tree is stored in the adjacency list $g$ g.

vector<vector<int> > g;

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

int dfs(int v, int p)
{

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

  sum[v] = 1;

Consider the edge $v \to t o$ v → to. If $t o \neq = p$ to \ne p, add to $s u m [v]$ sum[v] the number of vertices in the subtree rooted at $t o$ 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$ 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$ i, its immediate supervisor in the company is $x$ x.

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

Start the depth-first search from vertex $1$ 1.

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

Print the answer. The number of subordinates of employee $i$ i is equal to $s u m [i] -1$ 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

Python implementation

Increase the stack for recursion.

import sys
sys.setrecursionlimit(1000000)

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

def dfs(v, p):

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

  sum[v] = 1

Consider the edge $v \to t o$ v → to. If $t o \neq = p$ to \ne p, add to $s u m [v]$ sum[v] the number of vertices in the subtree rooted at $t o$ 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$ 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.

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 \geq 2)$ i~(i \ge 2), its immediate supervisor in the company is $a$ a.

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

Initialize the list $s u m$ sum.

sum = [0] * (n + 1)

Start the depth-first search from vertex $1$ 1.

dfs(1,-1)

Print the answer. The number of subordinates of employee $i$ i is equal to $s u m [i] - 1$ sum[i] - 1.

for i in range(1, n+1):
  print(sum[i] - 1,end=" ")

Eolymp #11428. Tree Diameter

A tree consisting of $n$ 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 \leq n \leq 2 \cdot 1 0^{5})$ n~(1 \le n \le 2 \cdot 10^5) — the number of vertices in the tree. The vertices are numbered from $1$ 1 to $n$ n.

Each of the following $n - 1$ n - 1 lines describes an edge and contains two integers $a$ a and $b (1 \leq a, b \leq n)$ b~(1 \le a, b \le n), indicating that there is an edge between vertices $a$ a and b.

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

Open problem

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 \le n \le 2 \cdot 10^5) — 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 \le a, b \le 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.

Open problem

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:]))

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 \le n \le 10^5) — 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 10^9.

Open problem

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++

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 \cdot (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 \cdot (n - c) \cdot w

    if (to != p)
    {
      int c = dfs(to, v);
      res = (res + 1LL * c * (n - c) * w) % MOD;
      cnt += c;
    }
  }
  return cnt;
}

C++

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++

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

Python implementation

Declare the value of the modulo as MOD.

MOD = 1000000000

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

c \cdot (n - c) \cdot 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

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 \le n \le 10^6). The following n lines describe the vertices of the tree. The i-th line contains two integers p_i and c_i, where p_i is the parent of vertex i, and c_i is the color of vertex i~(1 \le c_i \le n). For the root of the tree, p_i = 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.

Open problem

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++

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 \le n, c \le 10^5) — 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.

Open problem

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.

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++

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 \le n, q \le 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 p_i, q_i и r_i~(1 \le p_i, q_i \le n, 1 \le r_i \le 10^9), , meaning that videos p_i and q_i are connected by an edge with relevance r_i.

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

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.

Open problem

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++

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 \le n \le 10^5) 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.

Open problem

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++

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++

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 a_i and b_i. Vertex i is colored with color c_i (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 \le n \le 10^5) — the number of vertices.

The second line contains n integers — the colors of the vertices: c_1, c_2, \dots, c_n~(1 \le c_i \le 10^5).

Each of the following n - 1 lines contains two integers a_i and b_i~(1 \le a_i, b_i \le n) — the edges of the tree.

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

Open problem

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 \ne 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++

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 a_i and b_i.

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 10^9 + 7.

Input. The first line contains two numbers n and k~(1 \le n, k \le 10^5). Each of the following n - 1 lines contains two integers a_i and b_i~(1 \le a_i, b_i \le n).

Output. Print the number of ways to color the tree modulo 10^9 + 7

Open problem

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 = 10^9 + 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 \ne -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++

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 = 10^9 + 7. The parent of vertex v is p. Vertex v can be colored with col colors.

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

  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 10^9 + 7.

Input. The first line contains a single integer n~(2 \le n \le 2 \cdot 10^5) — the number of cities in the country. The second line contains n - 1 positive integers p_2, p_3, p_4, \cdots, p_n~(1 \le p_i \le i - 1), describing the roads in the country. The number p_i indicates that there is a road connecting city p_i and city i.

Output. Print the number of ways to improve the quality of the roads modulo 10^9 + 7.

Open problem

All calculations are performed modulo MOD = 10^9 + 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 to_1, to_2, ..., to_k are the children of v, then

f(v) = (f(to_1) + 1) \cdot (f(to_2) + 1) \cdot ... \cdot (f(to_k) + 1)

  for (int to : g[v])
    if (to != p)
    {
      long long sub = dfs(to, v);
      s = (s * (sub + 1)) % MOD;
    }
  return s;
}

C++

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++

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 \le n \le 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 \le u, v \le 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.

Open problem

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 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[to_1], h[to_2], ..., h[to_k]),

where to_1, to_2, ..., to_k 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[to_1], h[to_2], ..., h[to_k])

  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 to_1, to_2, ..., to_k be the children of vertex v. In the variable s, compute the sum:

res[to_1] + res[to_2] + ... + res[to_k]

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

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

res[v] = res[to_1] + res[to_2] + ... + res[to_k] = 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++

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 \le n \le 10^5, 1 \le m \le 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 a_1, a_2, ..., a_n, where each a_i 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 x_i~y_i~(1 \le x_i, y_i \le n, x_i \ne y_i), where x_i and y_i 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.

Open problem

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++

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

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(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 \le n \le 10^5) — 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.

Open problem

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++

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 \le n \le 3 \cdot 10^5).

Each of the following n - 1 lines contains two integers u_i and v_i~(1 \le u_i, v_i \le 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.

Open problem

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++

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++

Perform a depth-first search starting from vertex 1.

dfs(1);

Print the answer.

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

Trees

Fundamental Properties of a Tree

Alternative Definitions of a Tree

Representing Trees in Code

List of problems

Comments5