# Graph-coloring solution using depth-first search

I'm trying to solve this question:

Andryusha goes through a park each day. The squares and paths between them look boring to Andryusha, so he decided to decorate them.

The park consists of n squares connected with (n - 1) bidirectional paths in such a way that any square is reachable from any other using these paths. Andryusha decided to hang a colored balloon at each of the squares. The baloons' colors are described by positive integers, starting from 1. In order to make the park varicolored, Andryusha wants to choose the colors in a special way. More precisely, he wants to use such colors that if a, b and c are distinct squares that a and b have a direct path between them, and b and c have a direct path between them, then balloon colors on these three squares are distinct.

Andryusha wants to use as little different colors as possible. Help him to choose the colors!

Input The first line contains single integer n (3 ≤ n ≤ 2·105) — the number of squares in the park.

Each of the next (n - 1) lines contains two integers x and y (1 ≤ x, y ≤ n) — the indices of two squares directly connected by a path.

It is guaranteed that any square is reachable from any other using the paths.

Output In the first line print single integer k — the minimum number of colors Andryusha has to use.

In the second line print n integers, the i-th of them should be equal to the balloon color on the i-th square. Each of these numbers should be within range from 1 to k.

I have a solution that works, but times out on large inputs:

import java.util.*
import kotlin.collections.HashMap

private fun readInts() = readStrings().map { it.toInt() } // list of Ints

fun dfs(color_of: Array<Int>, graph: MutableMap<Int, MutableList<Int>>, current: Int, parent: Int) {
var c = 1
for (neb in graph[current] ?: mutableListOf()) {
if (color_of[neb] == 0) {
while (c == color_of[current] || c == color_of[parent]) {
c += 1
}
color_of[neb] = c
dfs(color_of, graph, neb, current)
c += 1
}
}
}

fun main(args: Array<String>) {
val graph = mutableMapOf<Int, MutableList<Int>>()
for (i in 0 until n - 1) {
graph.computeIfAbsent(x - 1) { mutableListOf() }
graph.computeIfAbsent(y - 1) { mutableListOf() }
}
val color_of = Array(n + 1) { 0 }
color_of = 1
dfs(color_of, graph, 0, 0)
println(color_of.max()!!)
for (i in 0 until n) {
println(color_of[i])
}
}



I'd love some pointers on how I can optimize this to pass big test cases.

Edit:

Here is the equivalent java solution that passes all test cases:

// Codeforces Round 403
// Andryusha and Colored Balloons
// Problem statement:
// http://codeforces.com/problemset/problem/781/A

import java.io.OutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.PrintWriter;
import java.io.PrintStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.StringTokenizer;
import java.io.InputStream;

public class Main {
public static void main(String[] args) {
InputStream inputStream = System.in;
OutputStream outputStream = System.out;
PrintWriter out = new PrintWriter(outputStream);
solver.solve(1, in, out);
out.close();
}

ArrayList<Integer>[] a;
int[] color;
int result;

public void solve(int testNumber, InputReader in, PrintWriter out) {
int n = in.nextInt();
color = new int[n];

a = new ArrayList[n];
for (int i = 0; i < n; i++) {
a[i] = new ArrayList<>();
}

for (int i = 0; i < n - 1; i++) {
int first = in.nextInt() - 1;
int second = in.nextInt() - 1;
}
color = 1;
result = 1;
dfs(0, 0);
out.println(result);
for (int i = 0; i < n; i++) {
out.print(color[i] + " ");
}
}

void dfs(int index, int parentIndex) {
// depth first traversal of graph
// paint each child node with color in ascending order skipping color of current node and of parent node
int currentColor = 1;
for (int childIndex : a[index]) {
if (color[childIndex] == 0) {
while ((currentColor == color[index]) || (currentColor == color[parentIndex])) currentColor++;
color[childIndex] = currentColor;
result = Math.max(result, currentColor);
dfs(childIndex, index);
currentColor++;
}
}
}

}

private static StringTokenizer tok;

}

public int nextInt() {
return Integer.parseInt(next());
}

public String next() {
try {
while (tok == null || !tok.hasMoreTokens()) {
}
} catch (IOException ex) {
System.err.println("An IOException was caught :" + ex.getMessage());
}