#include <iostream>
#include <vector>
#include <unordered_map>
#include <queue>
#include <algorithm>
using namespace std;
class Solution {
private :
vector< vector< int >> adj;
int n;
// BFS to find distance between two nodes
int bfs( int start, int end) {
if ( start == end) return 0 ;
vector< int > dist( n + 1 , - 1 ) ;
queue< int > q;
q.push ( start) ;
dist[ start] = 0 ;
while ( ! q.empty ( ) ) {
int u = q.front ( ) ;
q.pop ( ) ;
if ( u == end) return dist[ u] ;
for ( int v : adj[ u] ) {
if ( dist[ v] == - 1 ) {
dist[ v] = dist[ u] + 1 ;
q.push ( v) ;
}
}
}
return - 1 ;
}
public :
long long findTotalBandwidth( int network_nodes, vector< int > & network_from,
vector< int > & network_to, vector< int > & port) {
n = network_nodes;
adj.resize ( n + 1 ) ;
// Build adjacency list (1-indexed)
for ( int i = 0 ; i < network_from.size ( ) ; i++ ) {
int u = network_from[ i] ;
int v = network_to[ i] ;
adj[ u] .push_back ( v) ;
adj[ v] .push_back ( u) ;
}
// Group computers by port
unordered_map< int , vector< int >> portGroups;
for ( int i = 0 ; i < port.size ( ) ; i++ ) {
portGroups[ port[ i] ] .push_back ( i + 1 ) ; // 1-indexed
}
long long totalBandwidth = 0 ;
// For each port group, calculate sum of distances
for ( auto & [ p, computers] : portGroups) {
int groupSize = computers.size ( ) ;
// Calculate distance between all pairs in this group
for ( int i = 0 ; i < groupSize; i++ ) {
for ( int j = i + 1 ; j < groupSize; j++ ) {
int dist = bfs( computers[ i] , computers[ j] ) ;
totalBandwidth + = dist;
}
}
}
return totalBandwidth;
}
} ;
int main( ) {
Solution sol;
// Test Case 1: Sample Case 0
cout << "Test Case 1:" << endl;
{
int network_nodes = 5 ;
vector< int > network_from = { 1 , 1 , 2 , 4 } ;
vector< int > network_to = { 2 , 3 , 4 , 5 } ;
vector< int > port = { 1 , 2 , 2 , 3 , 2 } ;
long long result = sol.findTotalBandwidth ( network_nodes, network_from, network_to, port) ;
cout << "Expected: 8" << endl;
cout << "Got: " << result << endl;
cout << ( result == 8 ? "PASS" : "FAIL" ) << endl << endl;
}
// Test Case 2: Sample Case 1
cout << "Test Case 2:" << endl;
{
int network_nodes = 3 ;
vector< int > network_from = { 1 , 2 } ;
vector< int > network_to = { 2 , 3 } ;
vector< int > port = { 1 , 2 , 3 } ;
long long result = sol.findTotalBandwidth ( network_nodes, network_from, network_to, port) ;
cout << "Expected: 0" << endl;
cout << "Got: " << result << endl;
cout << ( result == 0 ? "PASS" : "FAIL" ) << endl << endl;
}
// Test Case 3: Example from problem description
cout << "Test Case 3:" << endl;
{
int network_nodes = 5 ;
vector< int > network_from = { 1 , 1 , 2 , 2 } ;
vector< int > network_to = { 2 , 3 , 4 , 5 } ;
vector< int > port = { 3 , 1 , 3 , 3 , 1 } ;
long long result = sol.findTotalBandwidth ( network_nodes, network_from, network_to, port) ;
cout << "Expected: 5" << endl;
cout << "Got: " << result << endl;
cout << ( result == 5 ? "PASS" : "FAIL" ) << endl << endl;
}
// Test Case 4: All same port
cout << "Test Case 4: All same port" << endl;
{
int network_nodes = 4 ;
vector< int > network_from = { 1 , 2 , 3 } ;
vector< int > network_to = { 2 , 3 , 4 } ;
vector< int > port = { 1 , 1 , 1 , 1 } ;
long long result = sol.findTotalBandwidth ( network_nodes, network_from, network_to, port) ;
// Distances: (1,2)=1, (1,3)=2, (1,4)=3, (2,3)=1, (2,4)=2, (3,4)=1
// Total = 1+2+3+1+2+1 = 10
cout << "Expected: 10" << endl;
cout << "Got: " << result << endl;
cout << ( result == 10 ? "PASS" : "FAIL" ) << endl << endl;
}
// Test Case 5: Star topology
cout << "Test Case 5: Star topology" << endl;
{
int network_nodes = 5 ;
vector< int > network_from = { 1 , 1 , 1 , 1 } ;
vector< int > network_to = { 2 , 3 , 4 , 5 } ;
vector< int > port = { 1 , 1 , 1 , 1 , 1 } ;
long long result = sol.findTotalBandwidth ( network_nodes, network_from, network_to, port) ;
// Center is node 1, distances from center: all 1
// Pairs: (1,2)=1, (1,3)=1, (1,4)=1, (1,5)=1
// Non-center pairs: (2,3)=2, (2,4)=2, (2,5)=2, (3,4)=2, (3,5)=2, (4,5)=2
// Total = 4*1 + 6*2 = 16
cout << "Expected: 16" << endl;
cout << "Got: " << result << endl;
cout << ( result == 16 ? "PASS" : "FAIL" ) << endl << endl;
}
// Test Case 6: Two separate groups
cout << "Test Case 6: Two groups" << endl;
{
int network_nodes = 4 ;
vector< int > network_from = { 1 , 2 , 3 } ;
vector< int > network_to = { 2 , 3 , 4 } ;
vector< int > port = { 1 , 1 , 2 , 2 } ;
long long result = sol.findTotalBandwidth ( network_nodes, network_from, network_to, port) ;
// Group 1: nodes 1,2 -> dist(1,2) = 1
// Group 2: nodes 3,4 -> dist(3,4) = 1
// Total = 1+1 = 2
cout << "Expected: 2" << endl;
cout << "Got: " << result << endl;
cout << ( result == 2 ? "PASS" : "FAIL" ) << endl << endl;
}
return 0 ;
}
I2luY2x1ZGUgPGlvc3RyZWFtPgojaW5jbHVkZSA8dmVjdG9yPgojaW5jbHVkZSA8dW5vcmRlcmVkX21hcD4KI2luY2x1ZGUgPHF1ZXVlPgojaW5jbHVkZSA8YWxnb3JpdGhtPgoKdXNpbmcgbmFtZXNwYWNlIHN0ZDsKCmNsYXNzIFNvbHV0aW9uIHsKcHJpdmF0ZToKICAgIHZlY3Rvcjx2ZWN0b3I8aW50Pj4gYWRqOwogICAgaW50IG47CiAgICAKICAgIC8vIEJGUyB0byBmaW5kIGRpc3RhbmNlIGJldHdlZW4gdHdvIG5vZGVzCiAgICBpbnQgYmZzKGludCBzdGFydCwgaW50IGVuZCkgewogICAgICAgIGlmIChzdGFydCA9PSBlbmQpIHJldHVybiAwOwogICAgICAgIAogICAgICAgIHZlY3RvcjxpbnQ+IGRpc3QobiArIDEsIC0xKTsKICAgICAgICBxdWV1ZTxpbnQ+IHE7CiAgICAgICAgcS5wdXNoKHN0YXJ0KTsKICAgICAgICBkaXN0W3N0YXJ0XSA9IDA7CiAgICAgICAgCiAgICAgICAgd2hpbGUgKCFxLmVtcHR5KCkpIHsKICAgICAgICAgICAgaW50IHUgPSBxLmZyb250KCk7CiAgICAgICAgICAgIHEucG9wKCk7CiAgICAgICAgICAgIAogICAgICAgICAgICBpZiAodSA9PSBlbmQpIHJldHVybiBkaXN0W3VdOwogICAgICAgICAgICAKICAgICAgICAgICAgZm9yIChpbnQgdiA6IGFkalt1XSkgewogICAgICAgICAgICAgICAgaWYgKGRpc3Rbdl0gPT0gLTEpIHsKICAgICAgICAgICAgICAgICAgICBkaXN0W3ZdID0gZGlzdFt1XSArIDE7CiAgICAgICAgICAgICAgICAgICAgcS5wdXNoKHYpOwogICAgICAgICAgICAgICAgfQogICAgICAgICAgICB9CiAgICAgICAgfQogICAgICAgIAogICAgICAgIHJldHVybiAtMTsKICAgIH0KICAgIApwdWJsaWM6CiAgICBsb25nIGxvbmcgZmluZFRvdGFsQmFuZHdpZHRoKGludCBuZXR3b3JrX25vZGVzLCB2ZWN0b3I8aW50PiYgbmV0d29ya19mcm9tLCAKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIHZlY3RvcjxpbnQ+JiBuZXR3b3JrX3RvLCB2ZWN0b3I8aW50PiYgcG9ydCkgewogICAgICAgIG4gPSBuZXR3b3JrX25vZGVzOwogICAgICAgIGFkai5yZXNpemUobiArIDEpOwogICAgICAgIAogICAgICAgIC8vIEJ1aWxkIGFkamFjZW5jeSBsaXN0ICgxLWluZGV4ZWQpCiAgICAgICAgZm9yIChpbnQgaSA9IDA7IGkgPCBuZXR3b3JrX2Zyb20uc2l6ZSgpOyBpKyspIHsKICAgICAgICAgICAgaW50IHUgPSBuZXR3b3JrX2Zyb21baV07CiAgICAgICAgICAgIGludCB2ID0gbmV0d29ya190b1tpXTsKICAgICAgICAgICAgYWRqW3VdLnB1c2hfYmFjayh2KTsKICAgICAgICAgICAgYWRqW3ZdLnB1c2hfYmFjayh1KTsKICAgICAgICB9CiAgICAgICAgCiAgICAgICAgLy8gR3JvdXAgY29tcHV0ZXJzIGJ5IHBvcnQKICAgICAgICB1bm9yZGVyZWRfbWFwPGludCwgdmVjdG9yPGludD4+IHBvcnRHcm91cHM7CiAgICAgICAgZm9yIChpbnQgaSA9IDA7IGkgPCBwb3J0LnNpemUoKTsgaSsrKSB7CiAgICAgICAgICAgIHBvcnRHcm91cHNbcG9ydFtpXV0ucHVzaF9iYWNrKGkgKyAxKTsgLy8gMS1pbmRleGVkCiAgICAgICAgfQogICAgICAgIAogICAgICAgIGxvbmcgbG9uZyB0b3RhbEJhbmR3aWR0aCA9IDA7CiAgICAgICAgCiAgICAgICAgLy8gRm9yIGVhY2ggcG9ydCBncm91cCwgY2FsY3VsYXRlIHN1bSBvZiBkaXN0YW5jZXMKICAgICAgICBmb3IgKGF1dG8mIFtwLCBjb21wdXRlcnNdIDogcG9ydEdyb3VwcykgewogICAgICAgICAgICBpbnQgZ3JvdXBTaXplID0gY29tcHV0ZXJzLnNpemUoKTsKICAgICAgICAgICAgCiAgICAgICAgICAgIC8vIENhbGN1bGF0ZSBkaXN0YW5jZSBiZXR3ZWVuIGFsbCBwYWlycyBpbiB0aGlzIGdyb3VwCiAgICAgICAgICAgIGZvciAoaW50IGkgPSAwOyBpIDwgZ3JvdXBTaXplOyBpKyspIHsKICAgICAgICAgICAgICAgIGZvciAoaW50IGogPSBpICsgMTsgaiA8IGdyb3VwU2l6ZTsgaisrKSB7CiAgICAgICAgICAgICAgICAgICAgaW50IGRpc3QgPSBiZnMoY29tcHV0ZXJzW2ldLCBjb21wdXRlcnNbal0pOwogICAgICAgICAgICAgICAgICAgIHRvdGFsQmFuZHdpZHRoICs9IGRpc3Q7CiAgICAgICAgICAgICAgICB9CiAgICAgICAgICAgIH0KICAgICAgICB9CiAgICAgICAgCiAgICAgICAgcmV0dXJuIHRvdGFsQmFuZHdpZHRoOwogICAgfQp9OwoKaW50IG1haW4oKSB7CiAgICBTb2x1dGlvbiBzb2w7CiAgICAKICAgIC8vIFRlc3QgQ2FzZSAxOiBTYW1wbGUgQ2FzZSAwCiAgICBjb3V0IDw8ICJUZXN0IENhc2UgMToiIDw8IGVuZGw7CiAgICB7CiAgICAgICAgaW50IG5ldHdvcmtfbm9kZXMgPSA1OwogICAgICAgIHZlY3RvcjxpbnQ+IG5ldHdvcmtfZnJvbSA9IHsxLCAxLCAyLCA0fTsKICAgICAgICB2ZWN0b3I8aW50PiBuZXR3b3JrX3RvID0gezIsIDMsIDQsIDV9OwogICAgICAgIHZlY3RvcjxpbnQ+IHBvcnQgPSB7MSwgMiwgMiwgMywgMn07CiAgICAgICAgCiAgICAgICAgbG9uZyBsb25nIHJlc3VsdCA9IHNvbC5maW5kVG90YWxCYW5kd2lkdGgobmV0d29ya19ub2RlcywgbmV0d29ya19mcm9tLCBuZXR3b3JrX3RvLCBwb3J0KTsKICAgICAgICBjb3V0IDw8ICJFeHBlY3RlZDogOCIgPDwgZW5kbDsKICAgICAgICBjb3V0IDw8ICJHb3Q6ICIgPDwgcmVzdWx0IDw8IGVuZGw7CiAgICAgICAgY291dCA8PCAocmVzdWx0ID09IDggPyAiUEFTUyIgOiAiRkFJTCIpIDw8IGVuZGwgPDwgZW5kbDsKICAgIH0KICAgIAogICAgLy8gVGVzdCBDYXNlIDI6IFNhbXBsZSBDYXNlIDEKICAgIGNvdXQgPDwgIlRlc3QgQ2FzZSAyOiIgPDwgZW5kbDsKICAgIHsKICAgICAgICBpbnQgbmV0d29ya19ub2RlcyA9IDM7CiAgICAgICAgdmVjdG9yPGludD4gbmV0d29ya19mcm9tID0gezEsIDJ9OwogICAgICAgIHZlY3RvcjxpbnQ+IG5ldHdvcmtfdG8gPSB7MiwgM307CiAgICAgICAgdmVjdG9yPGludD4gcG9ydCA9IHsxLCAyLCAzfTsKICAgICAgICAKICAgICAgICBsb25nIGxvbmcgcmVzdWx0ID0gc29sLmZpbmRUb3RhbEJhbmR3aWR0aChuZXR3b3JrX25vZGVzLCBuZXR3b3JrX2Zyb20sIG5ldHdvcmtfdG8sIHBvcnQpOwogICAgICAgIGNvdXQgPDwgIkV4cGVjdGVkOiAwIiA8PCBlbmRsOwogICAgICAgIGNvdXQgPDwgIkdvdDogIiA8PCByZXN1bHQgPDwgZW5kbDsKICAgICAgICBjb3V0IDw8IChyZXN1bHQgPT0gMCA/ICJQQVNTIiA6ICJGQUlMIikgPDwgZW5kbCA8PCBlbmRsOwogICAgfQogICAgCiAgICAvLyBUZXN0IENhc2UgMzogRXhhbXBsZSBmcm9tIHByb2JsZW0gZGVzY3JpcHRpb24KICAgIGNvdXQgPDwgIlRlc3QgQ2FzZSAzOiIgPDwgZW5kbDsKICAgIHsKICAgICAgICBpbnQgbmV0d29ya19ub2RlcyA9IDU7CiAgICAgICAgdmVjdG9yPGludD4gbmV0d29ya19mcm9tID0gezEsIDEsIDIsIDJ9OwogICAgICAgIHZlY3RvcjxpbnQ+IG5ldHdvcmtfdG8gPSB7MiwgMywgNCwgNX07CiAgICAgICAgdmVjdG9yPGludD4gcG9ydCA9IHszLCAxLCAzLCAzLCAxfTsKICAgICAgICAKICAgICAgICBsb25nIGxvbmcgcmVzdWx0ID0gc29sLmZpbmRUb3RhbEJhbmR3aWR0aChuZXR3b3JrX25vZGVzLCBuZXR3b3JrX2Zyb20sIG5ldHdvcmtfdG8sIHBvcnQpOwogICAgICAgIGNvdXQgPDwgIkV4cGVjdGVkOiA1IiA8PCBlbmRsOwogICAgICAgIGNvdXQgPDwgIkdvdDogIiA8PCByZXN1bHQgPDwgZW5kbDsKICAgICAgICBjb3V0IDw8IChyZXN1bHQgPT0gNSA/ICJQQVNTIiA6ICJGQUlMIikgPDwgZW5kbCA8PCBlbmRsOwogICAgfQogICAgCiAgICAvLyBUZXN0IENhc2UgNDogQWxsIHNhbWUgcG9ydAogICAgY291dCA8PCAiVGVzdCBDYXNlIDQ6IEFsbCBzYW1lIHBvcnQiIDw8IGVuZGw7CiAgICB7CiAgICAgICAgaW50IG5ldHdvcmtfbm9kZXMgPSA0OwogICAgICAgIHZlY3RvcjxpbnQ+IG5ldHdvcmtfZnJvbSA9IHsxLCAyLCAzfTsKICAgICAgICB2ZWN0b3I8aW50PiBuZXR3b3JrX3RvID0gezIsIDMsIDR9OwogICAgICAgIHZlY3RvcjxpbnQ+IHBvcnQgPSB7MSwgMSwgMSwgMX07CiAgICAgICAgCiAgICAgICAgbG9uZyBsb25nIHJlc3VsdCA9IHNvbC5maW5kVG90YWxCYW5kd2lkdGgobmV0d29ya19ub2RlcywgbmV0d29ya19mcm9tLCBuZXR3b3JrX3RvLCBwb3J0KTsKICAgICAgICAvLyBEaXN0YW5jZXM6ICgxLDIpPTEsICgxLDMpPTIsICgxLDQpPTMsICgyLDMpPTEsICgyLDQpPTIsICgzLDQpPTEKICAgICAgICAvLyBUb3RhbCA9IDErMiszKzErMisxID0gMTAKICAgICAgICBjb3V0IDw8ICJFeHBlY3RlZDogMTAiIDw8IGVuZGw7CiAgICAgICAgY291dCA8PCAiR290OiAiIDw8IHJlc3VsdCA8PCBlbmRsOwogICAgICAgIGNvdXQgPDwgKHJlc3VsdCA9PSAxMCA/ICJQQVNTIiA6ICJGQUlMIikgPDwgZW5kbCA8PCBlbmRsOwogICAgfQogICAgCiAgICAvLyBUZXN0IENhc2UgNTogU3RhciB0b3BvbG9neQogICAgY291dCA8PCAiVGVzdCBDYXNlIDU6IFN0YXIgdG9wb2xvZ3kiIDw8IGVuZGw7CiAgICB7CiAgICAgICAgaW50IG5ldHdvcmtfbm9kZXMgPSA1OwogICAgICAgIHZlY3RvcjxpbnQ+IG5ldHdvcmtfZnJvbSA9IHsxLCAxLCAxLCAxfTsKICAgICAgICB2ZWN0b3I8aW50PiBuZXR3b3JrX3RvID0gezIsIDMsIDQsIDV9OwogICAgICAgIHZlY3RvcjxpbnQ+IHBvcnQgPSB7MSwgMSwgMSwgMSwgMX07CiAgICAgICAgCiAgICAgICAgbG9uZyBsb25nIHJlc3VsdCA9IHNvbC5maW5kVG90YWxCYW5kd2lkdGgobmV0d29ya19ub2RlcywgbmV0d29ya19mcm9tLCBuZXR3b3JrX3RvLCBwb3J0KTsKICAgICAgICAvLyBDZW50ZXIgaXMgbm9kZSAxLCBkaXN0YW5jZXMgZnJvbSBjZW50ZXI6IGFsbCAxCiAgICAgICAgLy8gUGFpcnM6ICgxLDIpPTEsICgxLDMpPTEsICgxLDQpPTEsICgxLDUpPTEKICAgICAgICAvLyBOb24tY2VudGVyIHBhaXJzOiAoMiwzKT0yLCAoMiw0KT0yLCAoMiw1KT0yLCAoMyw0KT0yLCAoMyw1KT0yLCAoNCw1KT0yCiAgICAgICAgLy8gVG90YWwgPSA0KjEgKyA2KjIgPSAxNgogICAgICAgIGNvdXQgPDwgIkV4cGVjdGVkOiAxNiIgPDwgZW5kbDsKICAgICAgICBjb3V0IDw8ICJHb3Q6ICIgPDwgcmVzdWx0IDw8IGVuZGw7CiAgICAgICAgY291dCA8PCAocmVzdWx0ID09IDE2ID8gIlBBU1MiIDogIkZBSUwiKSA8PCBlbmRsIDw8IGVuZGw7CiAgICB9CiAgICAKICAgIC8vIFRlc3QgQ2FzZSA2OiBUd28gc2VwYXJhdGUgZ3JvdXBzCiAgICBjb3V0IDw8ICJUZXN0IENhc2UgNjogVHdvIGdyb3VwcyIgPDwgZW5kbDsKICAgIHsKICAgICAgICBpbnQgbmV0d29ya19ub2RlcyA9IDQ7CiAgICAgICAgdmVjdG9yPGludD4gbmV0d29ya19mcm9tID0gezEsIDIsIDN9OwogICAgICAgIHZlY3RvcjxpbnQ+IG5ldHdvcmtfdG8gPSB7MiwgMywgNH07CiAgICAgICAgdmVjdG9yPGludD4gcG9ydCA9IHsxLCAxLCAyLCAyfTsKICAgICAgICAKICAgICAgICBsb25nIGxvbmcgcmVzdWx0ID0gc29sLmZpbmRUb3RhbEJhbmR3aWR0aChuZXR3b3JrX25vZGVzLCBuZXR3b3JrX2Zyb20sIG5ldHdvcmtfdG8sIHBvcnQpOwogICAgICAgIC8vIEdyb3VwIDE6IG5vZGVzIDEsMiAtPiBkaXN0KDEsMikgPSAxCiAgICAgICAgLy8gR3JvdXAgMjogbm9kZXMgMyw0IC0+IGRpc3QoMyw0KSA9IDEKICAgICAgICAvLyBUb3RhbCA9IDErMSA9IDIKICAgICAgICBjb3V0IDw8ICJFeHBlY3RlZDogMiIgPDwgZW5kbDsKICAgICAgICBjb3V0IDw8ICJHb3Q6ICIgPDwgcmVzdWx0IDw8IGVuZGw7CiAgICAgICAgY291dCA8PCAocmVzdWx0ID09IDIgPyAiUEFTUyIgOiAiRkFJTCIpIDw8IGVuZGwgPDwgZW5kbDsKICAgIH0KICAgIAogICAgcmV0dXJuIDA7Cn0=