@@ -62,70 +62,67 @@ For every `i`, connect node `i` of the first tree with any node of the second tr
62
62
import kotlin.math.max
63
63
64
64
class Solution {
65
- private fun getGraph (edges : Array <IntArray >): Array <ArrayList <Int ?> ? > {
65
+ private fun getGraph (edges : Array <IntArray >): Array <ArrayList <Int > > {
66
66
val n = edges.size + 1
67
- val graph: Array <ArrayList <Int ?>? > = arrayOfNulls<ArrayList <Int ?>? > (n)
68
- for (i in 0 .. < n) {
69
- graph[i] = ArrayList <Int ?>()
70
- }
67
+ val graph: Array <ArrayList <Int >> = Array <ArrayList <Int >>(n) { ArrayList <Int >() }
71
68
for (edge in edges) {
72
69
val u = edge[0 ]
73
70
val v = edge[1 ]
74
- graph[u]!! .add(v)
75
- graph[v]!! .add(u)
71
+ graph[u].add(v)
72
+ graph[v].add(u)
76
73
}
77
74
return graph
78
75
}
79
76
80
- private fun dfs (graph : Array <ArrayList <Int ?> ? >, u : Int , pt : Int , dp : Array <IntArray ? >, k : Int ) {
81
- for (v in graph[u]!! ) {
77
+ private fun dfs (graph : Array <ArrayList <Int > >, u : Int , pt : Int , dp : Array <IntArray >, k : Int ) {
78
+ for (v in graph[u]) {
82
79
if (v == pt) {
83
80
continue
84
81
}
85
- dfs(graph, v!! , u, dp, k)
82
+ dfs(graph, v, u, dp, k)
86
83
for (i in 0 .. < k) {
87
- dp[u]!! [i + 1 ] + = dp[v]!! [i]
84
+ dp[u][i + 1 ] + = dp[v][i]
88
85
}
89
86
}
90
- dp[u]!! [0 ] = dp[u]!! [0 ] + 1
87
+ dp[u][0 ] = dp[u][0 ] + 1
91
88
}
92
89
93
90
private fun dfs2 (
94
- graph : Array <ArrayList <Int ?> ? >,
91
+ graph : Array <ArrayList <Int > >,
95
92
u : Int ,
96
93
pt : Int ,
97
94
ptv : IntArray ,
98
- fdp : Array <IntArray ? >,
99
- dp : Array <IntArray ? >,
95
+ fdp : Array <IntArray >,
96
+ dp : Array <IntArray >,
100
97
k : Int ,
101
98
) {
102
- fdp[u]!! [0 ] = dp[u]!! [0 ]
99
+ fdp[u][0 ] = dp[u][0 ]
103
100
for (i in 1 .. k) {
104
- fdp[u]!! [i] = (dp[u]!! [i] + ptv[i - 1 ])
101
+ fdp[u][i] = (dp[u][i] + ptv[i - 1 ])
105
102
}
106
- for (v in graph[u]!! ) {
103
+ for (v in graph[u]) {
107
104
if (v == pt) {
108
105
continue
109
106
}
110
107
val nptv = IntArray (k + 1 )
111
108
for (i in 0 .. < k) {
112
- nptv[i + 1 ] = dp[u]!! [i + 1 ] - dp[v!! ] !! [i] + ptv[i]
109
+ nptv[i + 1 ] = dp[u][i + 1 ] - dp[v] [i] + ptv[i]
113
110
}
114
111
nptv[0 ] = 1
115
- dfs2(graph, v!! , u, nptv, fdp, dp, k)
112
+ dfs2(graph, v, u, nptv, fdp, dp, k)
116
113
}
117
114
}
118
115
119
- private fun get (edges : Array <IntArray >, k : Int ): Array <IntArray ? > {
116
+ private fun get (edges : Array <IntArray >, k : Int ): Array <IntArray > {
120
117
val graph = getGraph(edges)
121
118
val n = graph.size
122
- val dp = Array <IntArray ? >(n) { IntArray (k + 1 ) }
123
- val fdp = Array <IntArray ? >(n) { IntArray (k + 1 ) }
119
+ val dp = Array <IntArray >(n) { IntArray (k + 1 ) }
120
+ val fdp = Array <IntArray >(n) { IntArray (k + 1 ) }
124
121
dfs(graph, 0 , - 1 , dp, k)
125
122
dfs2(graph, 0 , - 1 , IntArray (k + 1 ), fdp, dp, k)
126
123
for (i in 0 .. < n) {
127
124
for (j in 1 .. k) {
128
- fdp[i]!! [j] + = fdp[i]!! [j - 1 ]
125
+ fdp[i][j] + = fdp[i][j - 1 ]
129
126
}
130
127
}
131
128
return fdp
@@ -141,12 +138,12 @@ class Solution {
141
138
run {
142
139
var i = 0
143
140
while (k != 0 && i < m) {
144
- max = max(max.toDouble(), b[i]!! [k - 1 ].toDouble()).toInt()
141
+ max = max(max.toDouble(), b[i][k - 1 ].toDouble()).toInt()
145
142
i++
146
143
}
147
144
}
148
145
for (i in 0 .. < n) {
149
- ans[i] = a[i]!! [k] + max
146
+ ans[i] = a[i][k] + max
150
147
}
151
148
return ans
152
149
}
0 commit comments