설명#
💀 0~n-1까지 방향이 없는 트리와, 연결고리
edges가 주어졌을때,answer[i] = i번째 노드와 다른 모든 노드들과 거리의 합를 만족하는 answer을 리턴하기
def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
# n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]풀이과정#
1. 배열 2개 answer,node_cnt 생성#
answer : 정답용 answer[i] = 다른 모든 노드들과의 거리 합
node_cnt : 현재 정점의 누적 노드 개수
2. root를 기준으로 answer,node_cnt 채우기#
루트를 기준으로, 밑에서 위로올라오면서 배열 채우기
answer[i]= answer[i의 자식]+node_cnt[i의 자식] 의 합node_cnt[i]= i의 자식노드들 정점의 수 +1 (자기자신)
## dfs1()
//node_cnt
[6]#0 #root = 1+4+1=6
/ \
[1]#1 [4]#2 #1=1 #2=1+3 = 4
/ | \
[1][1][1]
//answer
[8]#0 #root = (0+1)+(3+4) =8
/ \
[0]#1 [3]#2 #1=> 0+(0*0) =0 #2=0+(3*1) = 3
/ | \
[0][0][0] #3,4,5=> 0+0 = 0 3.각 노드의 answer,node_cnt 채우기#
한번 순회했으면,
node_cnt는 완성된 상태고answer[0]의 값 (루트) 는 건들지 않아도 됨👍 하지만answer의다른 노드들은 밑에서 올라온 값만 계산한 상태임 따라서, 이번에는 위에서 내려가면서 배열을 갱신해줌**
answer[i]=(answer[i의부모] - node_cnt[i]) + (n-node_cnt[i]) ** 생각을 해보자 만일 2번 노드에 대한 정답을 구하려면, 루트노드와 2번의 하위 노드들을 토대로 구해야 한다
answer[0](루트)는 이미 2번에서 올라온 값을 가지고 있으므로, 이 값들은 중복됨. 따라서 node_cnt[i]만큼을 빼줘야 함- 반대로, 2번의 하위 노드가 아닌 얘들은, 길이가 하나씩 늘어날거임 => (전체 개수-node_cnt[2])를 하면 2번에 속하지 않은 애들 개수가 나오니, 이를 더해줌
## dfs2()
//node_cnt
[6]#0
/ \
[1]#1 [4]#2
/ | \
[1][1][1]
//answer (위에서부터 계산됨)
[8]#0 #root는 그대로
/ \
[12]#1 [6]#2 #1=> (8-1)+(6-1)= 12 #2=>(8-4)+(6-4)=6
/ | \
[5][5][5] #3,4,5=> (6-1)+(6-1) = 5 이러면 answer이 완성됨
전체 코드#
class Solution:
def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
tree = [[] for _ in range(n)]
answer = [0 for _ in range(n)]
node_cnt = [1 for _ in range(n)]
for a,b in edges:
tree[a].append(b)
tree[b].append(a)
def dfs(now, parent):
for child_node in tree[now]:
if child_node == parent:
continue;
dfs(child_node,now)
node_cnt[now] += node_cnt[child_node]
answer[now] += answer[child_node] + node_cnt[child_node]
dfs(0,-1)
def dfs2(now,parent):
for child_node in tree[now]:
if child_node == parent:
continue
answer[child_node] = (answer[now]-node_cnt[child_node]) + (n-node_cnt[child_node])
dfs2(child_node, now)
dfs2(0,-1)
return answer