Tree
Depth-First Search
Graph
Dynamic Programming
Description
There is an undirected connected tree with n
nodes labeled from 0
to n - 1
and n - 1
edges.
You are given the integer n
and the array edges
where edges[i] = [ai , bi ]
indicates that there is an edge between nodes ai
and bi
in the tree.
Return an array answer
of length n
where answer[i]
is the sum of the distances between the ith
node in the tree and all other nodes.
Example 1:
Input: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
Output: [8,12,6,10,10,10]
Explanation: The tree is shown above.
We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
equals 1 + 1 + 2 + 2 + 2 = 8.
Hence, answer[0] = 8, and so on.
Example 2:
Input: n = 1, edges = []
Output: [0]
Example 3:
Input: n = 2, edges = [[1,0]]
Output: [1,1]
Constraints:
1 <= n <= 3 * 104
edges.length == n - 1
edges[i].length == 2
0 <= ai , bi < n
ai != bi
The given input represents a valid tree.
Solutions
Solution 1
Python3 Java C++ Go TypeScript
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26 class Solution :
def sumOfDistancesInTree ( self , n : int , edges : List [ List [ int ]]) -> List [ int ]:
def dfs1 ( i : int , fa : int , d : int ):
ans [ 0 ] += d
size [ i ] = 1
for j in g [ i ]:
if j != fa :
dfs1 ( j , i , d + 1 )
size [ i ] += size [ j ]
def dfs2 ( i : int , fa : int , t : int ):
ans [ i ] = t
for j in g [ i ]:
if j != fa :
dfs2 ( j , i , t - size [ j ] + n - size [ j ])
g = defaultdict ( list )
for a , b in edges :
g [ a ] . append ( b )
g [ b ] . append ( a )
ans = [ 0 ] * n
size = [ 0 ] * n
dfs1 ( 0 , - 1 , 0 )
dfs2 ( 0 , - 1 , ans [ 0 ])
return ans
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42 class Solution {
private int n ;
private int [] ans ;
private int [] size ;
private List < Integer >[] g ;
public int [] sumOfDistancesInTree ( int n , int [][] edges ) {
this . n = n ;
g = new List [ n ] ;
ans = new int [ n ] ;
size = new int [ n ] ;
Arrays . setAll ( g , k -> new ArrayList <> ());
for ( var e : edges ) {
int a = e [ 0 ] , b = e [ 1 ] ;
g [ a ] . add ( b );
g [ b ] . add ( a );
}
dfs1 ( 0 , - 1 , 0 );
dfs2 ( 0 , - 1 , ans [ 0 ] );
return ans ;
}
private void dfs1 ( int i , int fa , int d ) {
ans [ 0 ] += d ;
size [ i ] = 1 ;
for ( int j : g [ i ] ) {
if ( j != fa ) {
dfs1 ( j , i , d + 1 );
size [ i ] += size [ j ] ;
}
}
}
private void dfs2 ( int i , int fa , int t ) {
ans [ i ] = t ;
for ( int j : g [ i ] ) {
if ( j != fa ) {
dfs2 ( j , i , t - size [ j ] + n - size [ j ] );
}
}
}
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37 class Solution {
public :
vector < int > sumOfDistancesInTree ( int n , vector < vector < int >>& edges ) {
vector < vector < int >> g ( n );
for ( auto & e : edges ) {
int a = e [ 0 ], b = e [ 1 ];
g [ a ]. push_back ( b );
g [ b ]. push_back ( a );
}
vector < int > ans ( n );
vector < int > size ( n );
function < void ( int , int , int ) > dfs1 = [ & ]( int i , int fa , int d ) {
ans [ 0 ] += d ;
size [ i ] = 1 ;
for ( int & j : g [ i ]) {
if ( j != fa ) {
dfs1 ( j , i , d + 1 );
size [ i ] += size [ j ];
}
}
};
function < void ( int , int , int ) > dfs2 = [ & ]( int i , int fa , int t ) {
ans [ i ] = t ;
for ( int & j : g [ i ]) {
if ( j != fa ) {
dfs2 ( j , i , t - size [ j ] + n - size [ j ]);
}
}
};
dfs1 ( 0 , -1 , 0 );
dfs2 ( 0 , -1 , ans [ 0 ]);
return ans ;
}
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33 func sumOfDistancesInTree ( n int , edges [][] int ) [] int {
g := make ([][] int , n )
for _ , e := range edges {
a , b := e [ 0 ], e [ 1 ]
g [ a ] = append ( g [ a ], b )
g [ b ] = append ( g [ b ], a )
}
ans := make ([] int , n )
size := make ([] int , n )
var dfs1 func ( i , fa , d int )
dfs1 = func ( i , fa , d int ) {
ans [ 0 ] += d
size [ i ] = 1
for _ , j := range g [ i ] {
if j != fa {
dfs1 ( j , i , d + 1 )
size [ i ] += size [ j ]
}
}
}
var dfs2 func ( i , fa , t int )
dfs2 = func ( i , fa , t int ) {
ans [ i ] = t
for _ , j := range g [ i ] {
if j != fa {
dfs2 ( j , i , t - size [ j ] + n - size [ j ])
}
}
}
dfs1 ( 0 , - 1 , 0 )
dfs2 ( 0 , - 1 , ans [ 0 ])
return ans
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30 function sumOfDistancesInTree ( n : number , edges : number [][]) : number [] {
const g : number [][] = Array . from ({ length : n }, () => []);
for ( const [ a , b ] of edges ) {
g [ a ]. push ( b );
g [ b ]. push ( a );
}
const ans : number [] = new Array ( n ). fill ( 0 );
const size : number [] = new Array ( n ). fill ( 0 );
const dfs1 = ( i : number , fa : number , d : number ) => {
ans [ 0 ] += d ;
size [ i ] = 1 ;
for ( const j of g [ i ]) {
if ( j !== fa ) {
dfs1 ( j , i , d + 1 );
size [ i ] += size [ j ];
}
}
};
const dfs2 = ( i : number , fa : number , t : number ) => {
ans [ i ] = t ;
for ( const j of g [ i ]) {
if ( j !== fa ) {
dfs2 ( j , i , t - size [ j ] + n - size [ j ]);
}
}
};
dfs1 ( 0 , - 1 , 0 );
dfs2 ( 0 , - 1 , ans [ 0 ]);
return ans ;
}
GitHub