My code runs fine while testing 1 Test case but fails while submit code

help
1

This is the problem:

Given an undirected and connected graph of N vertices and N-1 edges.
The number written on the ith vertex is val[i]. The task is to count
the number of pairs(x,y) such that the following conditions are satisfied:

  1. 0 ≤ x < y ≤ N-1
  2. The GCD of the numbers written on the vertices on the simple path in the given graph from x to y should be divisible by K.

For Example:
Input:
N = 4, K = 2
Edges[] = {{0, 1}, {0, 2}, {2, 3}}
Val[] = {2, 6, 4, 3}
Output:
3
Explanation:
0 - 1
|
2 - 3
There are three pairs - (0,1), (1,2) and
(0,2) which satisify both the conditions.

my solution:
import math

class Solution:
def countPair(self, N, Edges, Val, K):
count = 0

    for i in range(len(Val)):
        for j in range(len(Val) - 1, i, -1):
            if math.gcd(Val[i], Val[j]) % K == 0:
                count += 1
    return count

if name == ‘main’:
t = int(input())
for _ in range(t):
n,k = map(int, input().strip().split())
val = [int(x) for x in input().strip().split()]
edges = []
for i in range(n-1):
x,y = map(int, input().strip().split())
edges.append([x,y])
ob = Solution()
ans = ob.countPair(n,edges,val,k)
print(ans)

2

Please submit the problem link and a link to your code.
Or Format the code here so that its readable

3

Hey Roushan_17 - Here is the link - My code runs fine while testing 1 Test case but fails while submit code · GitHub

4

you can tag nodes as good or bad using rule (whether they are divisible by k or not)

then you can do union find and find the number of nodes in a good component(component in which all nodes are good). answer is sum of C(n, 2) for all good components.

This may also be solved using just dfs and some precomputation but haven’t tried that .

5

This problem is from the GFG contest, right?
My approach: For each connected component, such that every node is divisible by k in that component, do ans+=(len(comp)*(len(comp)-1)/2;

#define ll long long
class DisjointSet{
    unordered_map<ll,ll> parent,sz;
    public:
        void makeSet(vector<ll> &Nodes){
            for(ll i:Nodes){
                parent[i] = i;
                sz[i] = 1;
            }
        }    
        ll findParent(ll node){
            if(parent[node] ==  node)
                return node;
            return parent[node] = findParent(parent[node]);
        }
        void unionSet(ll a,ll b){
            a = findParent(a);
            b = findParent(b);
            if(a!=b){
                if(sz[a]<sz[b])
                    swap(a,b);
                parent[b]  = a;
                sz[a]+=sz[b];
            }
        }
};

long long countPairs(int N, vector<vector<int>> Edges, vector<int> val,int k)
    {
        // code here
        vector<ll> Nodes(N);
        for(ll i=0;i<N;i++)    Nodes[i] = i;
        DisjointSet djs;
        djs.makeSet(Nodes);
        for(int i=0;i<N-1;i++){
            int i1 = Edges[i][0],i2 = Edges[i][1];
            // cout<<i1<<" "<<i2;
            if((val[i1])%k == 0 && (val[i2])%k == 0){
                djs.unionSet(i1,i2);
            }
        }
        map<ll,ll> mp;
        ll ans = 0;
        for(int i=0;i<N;i++){
            //<<djs.findParent(i)<<" ";
            mp[djs.findParent(i)]++;
        }
        for(auto i:mp){
            ans+=((i.second)*(i.second-1))/2;
        }
        
        return ans;
    }