1
\$\begingroup\$

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 readLn() = readLine()!! // string line
private fun readInt() = readLn().toInt() // single Int
private fun readStrings() = readLn().trim().split(" ") // list of strings
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 n = readInt()
    val graph = mutableMapOf<Int, MutableList<Int>>()
    for (i in 0 until n - 1) {
        val (x, y) = readInts()
        graph.computeIfAbsent(x - 1) { mutableListOf() }
        graph.computeIfAbsent(y - 1) { mutableListOf() }
        graph[x - 1]!!.add(y - 1)
        graph[y - 1]!!.add(x - 1)
    }
    val color_of = Array(n + 1) { 0 }
    color_of[0] = 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.io.InputStreamReader;
import java.util.ArrayList;
import java.util.StringTokenizer;
import java.io.BufferedReader;
import java.io.InputStream;

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

    static class TaskC {
        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;
                a[first].add(second);
                a[second].add(first);
            }
            color[0] = 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++;
                }
            }
        }

    }

    static class InputReader {
        private static BufferedReader in;
        private static StringTokenizer tok;

        public InputReader(InputStream in) {
            this.in = new BufferedReader(new InputStreamReader(in));
        }

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

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

    }
}

What's wrong with MY solution? :(

\$\endgroup\$
  • 1
    \$\begingroup\$ when you searched the internet for this problem, how did others solve it? Did you already compare your solution to theirs? Especially for time-limit-exceeded problems, this helps in most cases. \$\endgroup\$ – Roland Illig Jul 5 at 21:14
  • \$\begingroup\$ @RolandIllig good point, I've added an equivalent java solution that passes all cases \$\endgroup\$ – nz_21 Jul 5 at 21:21
  • \$\begingroup\$ all code that you post here must be your own. Since you ask "what is wrong with MY solution", this implies that the second code snippet seems to be not yours. \$\endgroup\$ – Roland Illig Jul 6 at 6:22
  • \$\begingroup\$ What does profiling show to be the bottleneck? \$\endgroup\$ – Peter Taylor Aug 16 at 11:08

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.