fork download
copy 
  1. # factortree
  2.  
  3. #codechef factor tree problem
  4. #code
  5.  
  6. import math
  7.  
  8. mod=(10**9)+7
  9.  
  10. '''Function for calculating prime factors and adding the count of prime numbers in dictionary '''
  11.  
  12. def calc(path,a):
  13.  
  14.     #print(path)
  15.     dic={}
  16.     for i in range(len(path)):
  17.         dic=primeFactors(a[path[i]-1],dic)
  18.     value=list(dic.values())
  19.     ans=1
  20.     for i in range(len(value)):
  21.         ans*=(value[i]+1)
  22.         ans=ans%mod
  23.     print(ans)
  24.  
  25.  
  26. def primeFactors(n,dic): 
  27.  
  28.     while n % 2 == 0:
  29.         if 2 not in dic:
  30.             dic[2]=1
  31.         else:
  32.             dic[2]+=1
  33.         n = n / 2
  34.     for i in range(3,int(math.sqrt(n))+1,2): 
  35.         while n % i== 0: 
  36.             if i not in dic:
  37.                 dic[i]=1
  38.             else:
  39.                 dic[i]+=1 
  40.             n = n / i
  41.     if n > 2: 
  42.         if n not in dic:
  43.             dic[n]=1
  44.         else:
  45.             dic[n]+=1 
  46.     return dic
  47.  
  48. def addEdge(x, y,v): 
  49.  
  50.     v[x].append(y) 
  51.     v[y].append(x)  
  52.  
  53. def DFS(vis, x, y, stack,v,ansflag,a): 
  54.  
  55.     stack.append(x) 
  56.     if (x == y):
  57.         ansflag=1
  58.         calc(stack,a)
  59.         return 
  60.     vis[x] = True
  61.     flag = 0
  62.     if (len(v[x]) > 0): 
  63.         for j in v[x]: 
  64.             if (vis[j] == False): 
  65.                 DFS(vis, j, y, stack,v,ansflag,a) 
  66.                 if(ansflag==1):
  67.                     flag = 1  
  68.                     break
  69.  
  70.     if (flag == 0): 
  71.         del stack[-1] 
  72.  
  73.     return 
  74. def DFSCall(x, y, n, stack,v,a): 
  75.  
  76.     # visited array 
  77.     ansflag=0
  78.     vis = [0 for i in range(n + 1)] 
  79.     if(x!=y):
  80.         x1=DFS(vis, x, y, stack,v,ansflag,a)
  81.     else:
  82.         x1=[x]
  83.         calc(x1,a)
  84.     return
  85.  
  86. #main
  87.  
  88. for _ in range(int(input())):
  89.  
  90.     n = int(input())
  91.     v = [[] for i in range(n+1)]
  92.  
  93.     for i in range(n-1):
  94.         a1,b=map(int,input().split())
  95.         addEdge(a1,b,v) 
  96.     #print(v)
  97.     a=list(map(int,input().split()))
  98.     q=int(input())
  99.  
  100.     for i in range(q):
  101.         x,y=map(int,input().split())
  102.         stack = []
  103.         DFSCall(x,y,n,stack,v,a)
  104. #calling DFS for the path in tree
Success #stdin #stdout 0.02s 9232KB
comments (?)
stdin
copy 
1
5
1 2
1 3
2 4
2 5
2 6 4 3 5
2
1 4
2 2
stdout
copy 
9
4
https://ideone.com/2maMNH
language:
Python 3 (python  3.9.5)
created:
4 years ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!