Today, there was a test conducted by Uber and it's time slot was 8 AM to 2 PM (IST), so it has finished now. I need help in understanding the answer to a sample case of one of the problem. Problem was something like this:

There are N cities with a road between every pair of cities. Calculate the number of ways (modulo 1e9+7) such that you start travelling from **city 1** and return to the same **city 1** after completing **exactly K** trips where one trip means travelling from a current city to any other city.

Constraints: **N <= 1e9; K <= 1e6**

My approach: We start with city 1, now we have option to go (N-1) cities for 1st trip, then again (N-1) for 2nd trip, and so on. Hence, in short we have options of (N-1) cities for first (K-2) trips, then for the 2nd last trip we can go to (N-2) cities and for the last trip we have only 1 option that is to return to city 1. Hence according to this, answer should be: **(N-1)^(K-2) * (N-2)**

Sample Case: N=4, K=4

With my approach I am getting answer as **18**. But according to author the answer should be **21**. I need your help in understanding how there are 21 ways. I am not getting answer as 21 even when I try to find the ways using a pen and paper.

Appreciate your help. Thanks !

I guess the error is that the (K-1)th trip does not necessarily have (N-2) cities able to select. When the (K-2)th trip is at city 1, the (K-1)th trip would have (N-1) possible destinations.

Ahaan ! I don't know why I couldn't think of this. Now I get it where I was wrong! Thank you very much! I will try to formulate a solution according to this logic now

Is there any solution better than O(n*k)?

There is an O(k) solution.

Are u sure that N is as large as 1e9??

Yes

Wouldn't it be like follows?

from second trip to kth trip, there are N-2 choices only (all except start and current city).

Which makes the formula : (n-1)*(n-2)^(k-1)

It was nowhere mentioned that we couldn't return to city 1 in an intermediary trip, so we do have an option of going to city 1 at any trip as long as we end at city 1 after completing K trips, so there are N-1 choices at any stage

You can find the solution explanation here