Subtask 1

To get the points for subtask 1, it's sufficient to write a brute force solution that naively counts all possible valid paths. (See full solution below, but remove the parts regarding memoization and dynamic programming to make it exponentially slow.)

Subtask 2

Since the number of colors was at most ~3~, the length of the path was also at most ~3~; and thus either ~2~ or ~3~. Paths of length ~2~ are simple to count: they are just the number of edges between nodes with different colors. Paths of length ~3~ require a bit more thought. One way we can handle them is by looping over which node is in the middle of the path, and which colors the nodes at the start and end of the path have. Then the number of paths that have this node in the middle is the product of the number of neighbors of the first color and the number of neighbors of the second color.

Subtask 3

This is a natural extension to subtask 2: instead of looping over which node is in the middle of the path, we loop over which edge is in the middle. The rest works exactly the same.

Subtask 4

Assume there's a function ~f~ that gives the number of valid paths starting in a certain node. Then the answer to the problem is the sum of ~f~ over all nodes. The trick is to calculate ~f~ efficiently, as the input to the problem is quite large.

Let ~f~ take two parameters, ~c~ (the current node) and ~C~ (a bitset indicating the colors we have used so far, initially ~0~). Since the number of colors in a path is at most ~5~, the number of possible such bitsets is ~2^5 = 32~. This means that the number of possible combinations of parameters to ~f(c, C)~ is small enough to make a lookup table: we can use dynamic programming.

Memoizing on parameters ~c~ and ~C~, we implement ~f(c, C)~ by summing ~f(c', C')~ for all neighbors ~c'~ of ~c~, and where ~C'~ is the same bitset as ~C~ but with the color bit of ~c~ marked. We make sure to not take any paths where we reuse a color, and we make sure to not calculate answers for states that we have previously calculated. Given this, the answer to the problem is ~\sum f(i, 0)~ (make sure to not count paths of length ~1~!). Also remember to use 64-bit integer types.

This yields a time complexity of ~\mathcal O(2^K N)~.

Interesting note: the problem was essentially about counting paths of length at most ~k~, given that all nodes need to have different colors. This is much simpler than counting all paths of length at most ~k~! The naive time complexity of the latter is on the order of ~n^k~ rather than ~2^k n~. This can be used as a general algorithmic technique to speed up various problems: pick colors randomly a bunch of times (say, ~\mathcal O((\log n)^k)~ times), then use an algorithm that restricts nodes to have different colors.