Counting City Trips with K Stops

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++

// C++ Implementation
#include <bits/stdc++.h>
#include <iostream>
using namespace std;

int MOD = 1e9 + 7;
// Function to iterate over the roads
void dfs(vector<int> adj[], int u, int vis[], int& c)
{
    vis[u] = 1;
    c += 1;

    for (auto i : adj[u])

        // If city is not visited
        if (!vis[i])
            dfs(adj, i, vis, c);
}

// Function to calculate number of occurances possible
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 the number of occurances possible
    return ans;
}

// Function to calculate number of trips possible
int countTrips(int n, int k, vector<vector<int> > roads)
{

    vector<int> adj[n + 1];
    int vis[n + 1] = { 0 };

    // Form the adjancy matrix for normal roads
    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);
        }
    }

    // Calculate total combinations possible
    int ans = power(n, k);

    // Iterate over roads
    for (int i = 1; i <= n; i++) {
        int c = 0;
        // If node not visited
        if (!vis[i])
            dfs(adj, i, vis, c);

        // Add ans to variable
        ans = (ans - power(c, k) + MOD) % MOD;
    }

    // Return total trips possible
    return ans;
}

// Driver code
int main()
{

    int n = 5;
    int k = 3;

    vector<vector<int> > roads = {
        { 1, 2, 1 }, { 2, 3, 1 }, { 3, 4, 1 }, { 4, 5, 0 }
    };

    // Function call
    cout << countTrips(n, k, roads);
    return 0;
}

Java

import java.util.ArrayList;
import java.util.List;

public class CountTrips {

    static int MOD = 1000000007;

    // Function to iterate over the roads
    static void dfs(List<Integer>[] adj, int u, int[] vis, int[] c) {
        vis[u] = 1;
        c[0] += 1;

        for (int i : adj[u]) {
            // If city is not visited
            if (vis[i] == 0) {
                dfs(adj, i, vis, c);
            }
        }
    }

    // Function to calculate the number of occurrences possible
    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 the number of occurrences possible
        return ans;
    }

    // Function to calculate the number of trips possible
    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];

        // Form the adjacency list for normal roads
        for (int[] it : roads) {
            int u = it[0], v = it[1];
            if (it[2] == 0) {
                adj[u].add(v);
                adj[v].add(u);
            }
        }

        // Calculate total combinations possible
        int ans = (int) power(n, k);

        // Iterate over roads
        for (int i = 1; i <= n; i++) {
            int[] c = { 0 };
            // If node not visited
            if (vis[i] == 0)
                dfs(adj, i, vis, c);

            // Add ans to variable
            ans = (int) ((ans - power(c[0], k) + MOD) % MOD);
        }

        // Return total trips possible
        return ans;
    }

    // Driver code
    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});

        // Function call
        System.out.println(countTrips(n, k, roads));
    }
}


// This code is contributed by rambabuguphka

Python3

from collections import defaultdict

MOD = 1000000007

# Function to iterate over the roads
def dfs(adj, u, vis, c):
    vis[u] = 1
    c[0] += 1

    for i in adj[u]:
        # If city is not visited
        if vis[i] == 0:
            dfs(adj, i, vis, c)

# Function to calculate the number of occurrences possible
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 the number of occurrences possible
    return ans

# Function to calculate the number of trips possible
def countTrips(n, k, roads):
    adj = defaultdict(list)
    vis = [0] * (n + 1)

    # Form the adjacency list for normal roads
    for u, v, road_type in roads:
        if road_type == 0:
            adj[u].append(v)
            adj[v].append(u)

    # Calculate total combinations possible
    ans = power(n, k)

    # Iterate over roads
    for i in range(1, n + 1):
        c = [0]
        # If node not visited
        if vis[i] == 0:
            dfs(adj, i, vis, c)

        # Add ans to variable
        ans = (ans - power(c[0], k) + MOD) % MOD

    # Return total trips possible
    return ans

# Driver code
if __name__ == "__main__":
    n = 5
    k = 3

    roads = [
        [1, 2, 1],
        [2, 3, 1],
        [3, 4, 1],
        [4, 5, 0]
    ]

    # Function call
    print(countTrips(n, k, roads))

C#

using System;
using System.Collections.Generic;

class Program
{
    static int MOD = 1000000007;

    // Function to iterate over the roads
    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 city is not visited
            if (vis[i] == 0)
                Dfs(adj, i, vis, ref c);
        }
    }

    // Function to calculate power
    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;
    }

    // Function to calculate number of trips possible
    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;
    }

    // Driver code
    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));
    }
}

// This code is contributed by shivamgupta0987654321

Javascript

// Javascript Implementation

const MOD = 1e9 + 7;

// Function to perform depth-first search on roads
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 to calculate power in modular arithmetic
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 to calculate the number of trips possible
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;
}

// Driver code

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));


// This code is contributed by Tapesh(tapeshdua420)

Output

114

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).