Counting City Trips with K Stops
Last Updated :
18 Jan, 2024
Given n cities connected by (n-1) roads, each specified by [a, b, type], denoting a connection between cities a and b of a certain ‘type’ (0 for normal roads, 1 for highways). The task is to count the number of trips satisfying two conditions:
- The trip must include travel via a highway at least once.
- The trip must have exactly K city stops along the way.
Return the final result by taking modulo 10^9 + 7. Additionally, it’s emphasized that stops during the trip need not be distinct.
Examples:
Input: n = 3, k = 2, roads = {{1, 2, 1}, {2, 3, 0}}
Output: 4
Explanation: Possible good trips are:
{1, 2}: start at 1, stop at 2.
{1, 3}: start at 1, don’t stop at 2, go directly to 3.
{2, 1}: start at 2, stop at 1.
{3, 1}: start at 3, don’t stop at 2, go directly to 1.
Input: n = 5, k = 3, roads = { {1, 2, 1}, {2, 3, 1}, {3, 4, 1}, {4, 5, 0} }
Output: 114
Explanation: All trips are good except – {1, 1, 1}, {2, 2, 2}, {3, 3, 3}, {4, 4, 4}, {5, 5, 5}, {4, 4, 5}, {4, 5, 4}, {4, 5, 5}, {5, 4, 4}, {5, 4, 5}, {5, 5, 4} as they do not involve travelling via a highway.
Approach: This can be solved with the following idea:
Count total number of trips that can be possible by avoiding above mentioned conditions. Then calculate trips made by normal road. Subtract both of them which will led to our desired ans.
Step-by-step approach:
- Create a adjacency matrix adj for normal roads.
- Calculate total combination possible by using power function.
- Iterate over normal roads:
- As it will help us to find out atleast one highway combinations.
- Using DFS function, we can find combination possible.
- Add combination to ans.
- Return ans.
Below is the implementation of the above approach:
C++
#include <bits/stdc++.h>
#include <iostream>
using namespace std;
int MOD = 1e9 + 7;
void dfs(vector< int > adj[], int u, int vis[], int & c)
{
vis[u] = 1;
c += 1;
for ( auto i : adj[u])
if (!vis[i])
dfs(adj, i, vis, c);
}
long power( long a, long b)
{
if (b == 0)
return 1;
long ans = 1;
while (b > 0) {
if (b & 1)
ans = (ans * a) % MOD;
a = (a * a) % MOD;
b >>= 1;
}
return ans;
}
int countTrips( int n, int k, vector<vector< int > > roads)
{
vector< int > adj[n + 1];
int vis[n + 1] = { 0 };
for ( auto it : roads) {
int u = it[0], v = it[1];
if (it[2] == 0) {
adj[u].push_back(v);
adj[v].push_back(u);
}
}
int ans = power(n, k);
for ( int i = 1; i <= n; i++) {
int c = 0;
if (!vis[i])
dfs(adj, i, vis, c);
ans = (ans - power(c, k) + MOD) % MOD;
}
return ans;
}
int main()
{
int n = 5;
int k = 3;
vector<vector< int > > roads = {
{ 1, 2, 1 }, { 2, 3, 1 }, { 3, 4, 1 }, { 4, 5, 0 }
};
cout << countTrips(n, k, roads);
return 0;
}
|
Java
import java.util.ArrayList;
import java.util.List;
public class CountTrips {
static int MOD = 1000000007 ;
static void dfs(List<Integer>[] adj, int u, int [] vis, int [] c) {
vis[u] = 1 ;
c[ 0 ] += 1 ;
for ( int i : adj[u]) {
if (vis[i] == 0 ) {
dfs(adj, i, vis, c);
}
}
}
static long power( long a, long b) {
if (b == 0 )
return 1 ;
long ans = 1 ;
while (b > 0 ) {
if ((b & 1 ) == 1 )
ans = (ans * a) % MOD;
a = (a * a) % MOD;
b >>= 1 ;
}
return ans;
}
static int countTrips( int n, int k, List< int []> roads) {
List<Integer>[] adj = new ArrayList[n + 1 ];
for ( int i = 0 ; i <= n; i++) {
adj[i] = new ArrayList<>();
}
int [] vis = new int [n + 1 ];
for ( int [] it : roads) {
int u = it[ 0 ], v = it[ 1 ];
if (it[ 2 ] == 0 ) {
adj[u].add(v);
adj[v].add(u);
}
}
int ans = ( int ) power(n, k);
for ( int i = 1 ; i <= n; i++) {
int [] c = { 0 };
if (vis[i] == 0 )
dfs(adj, i, vis, c);
ans = ( int ) ((ans - power(c[ 0 ], k) + MOD) % MOD);
}
return ans;
}
public static void main(String[] args) {
int n = 5 ;
int k = 3 ;
List< int []> roads = new ArrayList<>();
roads.add( new int []{ 1 , 2 , 1 });
roads.add( new int []{ 2 , 3 , 1 });
roads.add( new int []{ 3 , 4 , 1 });
roads.add( new int []{ 4 , 5 , 0 });
System.out.println(countTrips(n, k, roads));
}
}
|
Python3
from collections import defaultdict
MOD = 1000000007
def dfs(adj, u, vis, c):
vis[u] = 1
c[ 0 ] + = 1
for i in adj[u]:
if vis[i] = = 0 :
dfs(adj, i, vis, c)
def power(a, b):
if b = = 0 :
return 1
ans = 1
while b > 0 :
if b & 1 = = 1 :
ans = (ans * a) % MOD
a = (a * a) % MOD
b >> = 1
return ans
def countTrips(n, k, roads):
adj = defaultdict( list )
vis = [ 0 ] * (n + 1 )
for u, v, road_type in roads:
if road_type = = 0 :
adj[u].append(v)
adj[v].append(u)
ans = power(n, k)
for i in range ( 1 , n + 1 ):
c = [ 0 ]
if vis[i] = = 0 :
dfs(adj, i, vis, c)
ans = (ans - power(c[ 0 ], k) + MOD) % MOD
return ans
if __name__ = = "__main__" :
n = 5
k = 3
roads = [
[ 1 , 2 , 1 ],
[ 2 , 3 , 1 ],
[ 3 , 4 , 1 ],
[ 4 , 5 , 0 ]
]
print (countTrips(n, k, roads))
|
C#
using System;
using System.Collections.Generic;
class Program
{
static int MOD = 1000000007;
static void Dfs(List< int >[] adj, int u, int [] vis, ref int c)
{
vis[u] = 1;
c += 1;
foreach ( var i in adj[u])
{
if (vis[i] == 0)
Dfs(adj, i, vis, ref c);
}
}
static long Power( long a, long b)
{
if (b == 0)
return 1;
long ans = 1;
while (b > 0)
{
if ((b & 1) == 1)
ans = (ans * a) % MOD;
a = (a * a) % MOD;
b >>= 1;
}
return ans;
}
static int CountTrips( int n, int k, List<List< int >> roads)
{
List< int >[] adj = new List< int >[n + 1];
for ( int i = 0; i <= n; i++)
adj[i] = new List< int >();
int [] vis = new int [n + 1];
foreach ( var it in roads)
{
int u = it[0], v = it[1];
if (it[2] == 0)
{
adj[u].Add(v);
adj[v].Add(u);
}
}
long ans = Power(n, k);
for ( int i = 1; i <= n; i++)
{
int c = 0;
if (vis[i] == 0)
Dfs(adj, i, vis, ref c);
ans = (ans - Power(c, k) + MOD) % MOD;
}
return ( int )ans;
}
static void Main( string [] args)
{
int n = 5;
int k = 3;
List<List< int >> roads = new List<List< int >>
{
new List< int > { 1, 2, 1 },
new List< int > { 2, 3, 1 },
new List< int > { 3, 4, 1 },
new List< int > { 4, 5, 0 }
};
Console.WriteLine(CountTrips(n, k, roads));
}
}
|
Javascript
const MOD = 1e9 + 7;
function dfs(adj, u, vis) {
vis[u] = 1;
let c = 1;
for (let i of adj[u]) {
if (!vis[i]) {
c += dfs(adj, i, vis);
}
}
return c;
}
function power(a, b) {
if (b === 0) return 1;
let ans = 1;
while (b > 0) {
if (b & 1) {
ans = (ans * a) % MOD;
}
a = (a * a) % MOD;
b >>= 1;
}
return ans;
}
function countTrips(n, k, roads) {
const adj = Array.from({ length: n + 1 }, () => []);
const vis = Array(n + 1).fill(0);
for (let it of roads) {
const [u, v, flag] = it;
if (flag === 0) {
adj[u].push(v);
adj[v].push(u);
}
}
let ans = power(n, k);
for (let i = 1; i <= n; i++) {
if (!vis[i]) {
const c = dfs(adj, i, vis);
ans = (ans - power(c, k) + MOD) % MOD;
}
}
return ans;
}
const n = 5;
const k = 3;
const roads = [
[1, 2, 1],
[2, 3, 1],
[3, 4, 1],
[4, 5, 0]
];
console.log(countTrips(n, k, roads));
|
Time Complexity: O(V + E), where V is the number of vertices (cities) and E is the number of edges (roads).
Auxiliary Space: O(V), where V is the number of vertices (cities).
Share your thoughts in the comments
Please Login to comment...