본문으로 건너뛰기

3372 - Maximize the Number of Target Nodes After Connecting Trees I

info

풀이 키워드​

스포주의

dfs


풀이 코드​

info
  • 메모리: 46760 KB
  • 시간: 279 ms
class Solution {
    List<Integer>[] adj1, adj2;
    boolean[] vis1, vis2;
    int n, m, k;

    public int dfs1(int i, int depth) { // tree 1
        if (depth > k) return 0;

        int res = 1;
        for (int nxt : adj1[i]) {
            vis1[i] = true;
            if (!vis1[nxt]) res += dfs1(nxt, depth+1);
            vis1[i] = false;
        }

        return res;
    }

    public int dfs2(int j, int depth) { // tree 2
        if (depth > k-1) return 0;

        int res = 1;
        for (int nxt : adj2[j]) {
            vis2[j] = true;
            if (!vis2[nxt]) res += dfs2(nxt, depth+1);
            vis2[j] = false;
        }

        return res;
    }

    public int[] maxTargetNodes(int[][] edges1, int[][] edges2, int k) {
        this.n = edges1.length + 1;
        this.m = edges2.length + 1;
        this.k = k;

        // init adj
        this.adj1 = new ArrayList[n];
        this.adj2 = new ArrayList[m];
        for (int i = 0; i < n; ++i) adj1[i] = new ArrayList<>();
        for (int j = 0; j < m; ++j) adj2[j] = new ArrayList<>();
        for (int[] e : edges1) {
            adj1[e[0]].add(e[1]);
            adj1[e[1]].add(e[0]);
        }
        for (int[] e : edges2) {
            adj2[e[0]].add(e[1]);
            adj2[e[1]].add(e[0]);
        }

        // init visited
        vis1 = new boolean[n];
        vis2 = new boolean[m];

        // dfs
        int[] ans = new int[n];
        int mx = 0; // get maximum tree 2 nodes of distance k
        for (int j = 0; j < m; ++j)
            mx = Math.max(mx, dfs2(j, 0));
        
        for (int i = 0; i < n; ++i) {
            // count tree 1 nodes of distance k
            ans[i] = dfs1(i, 0) + mx;
        }

        return ans;
    }
}

풀이 해설​

tree1의 특정 노드 i를 tree2의 노드에 연결했을 때, i로부터의 거리가 k 이하인 최대 노드 수를 구하는 문제이다.

정확히는 tree1의 모든 노드에 대하여 구해야되기 때문에 반환형이 int가 아니고 int[]이다.

지문이 좀 난독증 오는데 댓글에서 사람들이 남긴 설명을 읽어보면 그제서야 이해된다.


By maximum possible number of nodes target to node i of the first tree, it means how many nodes are less than k edges away from node i of the first tree.

For Example 1: edges1 = [[0,1],[0,2],[2,3],[2,4]], edges2 = [[0,1],[0,2],[0,3],[2,7],[1,4],[4,5],[4,6]], k = 2

For i = 0, connect node 0 from the first tree to node 0 from the second tree.

If there is an edge from node 0 of tree1 and node 0 of tree2
Here are the nodes k edges away to node 0 of tree1:
from tree1: [0,1,2,3,4]
from tree2: [0,2,3,1]
For a total of 9 nodes


😮 tree2 dfs는 한번만 돌려...​

tree1의 모든 노드에 대해 구하라고 되어있어서 tree1의 모든 노드(i:[0,n)i: [0, n))에 대한 순회 로직이 최외각에 있었고

이 순회 로직 내부에서 tree2에 대한 dfs2를 호출하는 실수를 했는데, 당연히 TLE 터졌다.

생각해보니 dfs2는 tree2에서의 max 값 찾는 목적으로 돌리는 것이라, 0 ~ m-1 노드까지 한번 돌려놓으면 끝이다.

이후는 그냥 각 노드 i 별로 max 값만 합해서 정답 배열에 넣어주면 되기 때문이다.

그래서 dfs1 돌리기 전에 먼저 수행하는 것으로 수정했다.


🚫 무지성 base case 작성 금지​

WA1
public int dfs2(int j, int depth) { // tree 2
    if (depth == k-1) return 1;
WA2
public int dfs2(int j, int depth) { // tree 2
    if (depth >= k-1) return 1;

원래는 상단과 같은 EQ, GTE로 기저 조건을 작성했는데 다 WA 받았다.

WA가 뜨는 TC는 다음과 같다.

[[0,1]]
[[0,1]]
0
Expected: [1,1]

WA를 받은 이유는 k가 0일 수 있기 때문이다.

  • 0 <= k <= 1000

문제를 잘 읽자

depth == -1 는 기저 조건이 작동을 안하고

depth >= -1 은 작동하지만 return 1을 하므로 틀렸다. 애초에 시작부터 거리 초과라서 노드 j를 카운팅하면 안되기 때문이다.

AC
public int dfs2(int j, int depth) { // tree 2
    if (depth > k-1) return 0;

따라서 거리가 초과되었을 때 0을 리턴하도록 하는 것이 옳다.


✨ 번외: TC 시각화하기​

811번 효율성 저격용 테스트케이스를 파이썬으로 시각화해보았다.

코드가 가성비로 맛있음ㅋㅋㅋ

import networkx as nx
import matplotlib.pyplot as plt

edges1 = [[257,1],[340,3],[454,5],[847,6],[701,10],[436,11],[710,16],[14,17],[495,14],[259,20],[858,21],[849,28],[733,30],[780,35],[304,43],[567,46],[851,47],[409,48],[606,52],[466,54],[78,57],[18,58],[537,62],[786,65],[482,69],[770,70],[132,71],[31,72],[183,31],[699,73],[27,74],[7,77],[157,78],[97,82],[744,83],[812,84],[496,88],[504,92],[600,93],[153,94],[714,96],[299,98],[139,99],[296,101],[441,103],[784,104],[283,105],[39,109],[670,110],[391,114],[435,115],[90,116],[435,118],[317,124],[725,125],[209,126],[134,127],[681,130],[59,131],[209,132],[227,137],[36,139],[444,144],[319,145],[341,147],[141,149],[819,155],[203,156],[87,159],[795,87],[79,168],[750,174],[468,175],[409,177],[310,179],[690,180],[658,181],[292,182],[722,183],[426,186],[321,187],[258,188],[646,192],[148,195],[97,200],[472,97],[277,203],[257,204],[595,206],[59,208],[568,59],[306,211],[36,217],[517,36],[18,220],[292,222],[251,223],[44,225],[198,44],[569,198],[651,229],[407,231],[7,232],[113,233],[112,236],[13,239],[774,244],[527,245],[697,246],[768,247],[828,251],[237,253],[769,237],[475,254],[604,256],[573,258],[846,263],[407,267],[378,272],[13,278],[637,279],[45,280],[347,282],[294,288],[597,290],[768,292],[573,295],[815,296],[75,301],[431,75],[360,303],[530,306],[609,308],[444,310],[330,314],[51,318],[283,319],[716,324],[829,329],[294,331],[150,294],[162,333],[298,336],[729,298],[261,340],[800,261],[413,342],[735,345],[215,346],[459,215],[432,349],[506,351],[719,352],[563,357],[341,360],[720,363],[844,364],[644,366],[471,367],[559,372],[228,373],[634,228],[496,374],[670,376],[723,378],[509,379],[270,384],[640,270],[285,385],[334,285],[230,387],[572,390],[154,394],[100,154],[542,100],[795,396],[135,397],[547,135],[240,399],[320,240],[45,401],[783,45],[120,403],[447,120],[835,411],[219,412],[571,413],[639,414],[698,415],[767,420],[485,422],[434,423],[647,425],[128,430],[710,431],[242,433],[350,435],[671,350],[497,436],[300,438],[753,442],[37,443],[520,37],[493,446],[106,447],[459,448],[34,450],[453,34],[128,453],[210,128],[347,454],[375,347],[328,457],[441,459],[514,460],[209,461],[816,209],[7,462],[445,7],[238,463],[51,464],[317,51],[798,465],[249,469],[771,471],[382,472],[482,474],[577,478],[164,479],[391,164],[334,480],[171,481],[793,485],[714,487],[289,488],[632,490],[23,492],[138,23],[539,494],[238,499],[737,238]]
edges2 = [[645,0],[690,1],[612,4],[561,8],[111,11],[646,12],[718,13],[747,14],[10,16],[546,20],[109,21],[143,27],[408,31],[278,34],[674,35],[453,38],[28,46],[355,47],[596,48],[559,50],[39,51],[368,53],[83,56],[642,61],[446,62],[191,64],[435,65],[155,67],[723,69],[166,76],[415,80],[673,85],[528,87],[25,90],[97,25],[691,96],[282,97],[681,98],[338,99],[518,104],[108,105],[411,107],[313,109],[699,112],[111,118],[416,120],[45,121],[334,123],[722,132],[704,140],[163,142],[333,144],[552,145],[470,147],[266,149],[698,150],[348,154],[641,157],[384,158],[88,162],[340,163],[182,166],[233,169],[184,171],[70,173],[83,70],[405,83],[517,174],[280,178],[302,180],[381,181],[720,183],[317,187],[602,191],[453,192],[282,193],[117,195],[718,196],[651,198],[82,200],[271,203],[754,204],[359,206],[398,211],[654,214],[86,216],[254,86],[202,221],[631,202],[141,223],[289,141],[15,224],[352,226],[630,234],[284,236],[448,237],[542,240],[201,241],[372,201],[495,245],[408,246],[382,247],[371,248],[529,252],[525,257],[422,259],[42,261],[100,42],[750,263],[253,264],[40,253],[190,40],[379,190],[306,266],[469,271],[45,273],[477,277],[332,279],[320,280],[512,283],[146,284],[289,285],[127,289],[418,127]]

# Combine all edges
all_edges = edges1 + edges2

# Create a graph
G = nx.Graph()
G.add_edges_from(all_edges)

# Draw the graph
plt.figure(figsize=(16, 14))
pos = nx.spring_layout(G, k=0.15)

node_colors = ['pink' if node == 0 else 'skyblue' for node in G.nodes()]
nx.draw_networkx_nodes(G, pos, node_size=50, node_color=node_colors)
nx.draw_networkx_labels(G, pos, font_size=8)

nx.draw_networkx_edges(G, pos, alpha=0.3)

plt.title("Graph Representation from Combined Edge Lists", fontsize=20)
plt.axis("off")
plt.show()

https://velog.velcdn.com/images/qriosity/post/7727006f-85d4-46bd-aaf5-8d3032ba5204/image.png

메모​

  • 구현은 쉬운데, 문제를 필요 이상으로 어렵게 해석하지 않도록 주의
    • 단순히 dfs 두 개 쓰면 되는걸 또 이상한 DP 의심하고 있었음 ㅜ