Openjudge 百练第4109题

在OpenJudge看到一个题目（#4109），题目描述如下：

小明和小红去参加party。会场中总共有n个人，这些人中有的是朋友关系，有的则相互不认识。朋友关系是相互的，即如果A是B的朋友，那么B也是A的朋友。小明和小红想知道其中某两个人有多少个公共的朋友。

输入第一行为一个正整数c，代表测试数据的个数。接下来是c组测试数据。

对于每组测试数据，第一行是三个数字n(2<=n<=100)，m和k，分别表示会场中的人数，已知的朋友关系数目，问题的数目。接下来的m行，每行用两个数字i和j(1<=i,j<=n)表示了一个朋友关系，表示第i个人和第j个人是朋友关系。接下来的k行，每行用两个数字i和j(1<=i,j<=n)表示一个问题，请问第i个人和第j个人有多少公共的朋友。输出对于第i组测试数据，首先输出一行”Case i:”，接下来得k行代表了k个问题，每行输出第i个人和第j个人有多少公共的朋友。

用Python写了段代码，大致实现：

  # -*- coding:utf-8 -*-
 
  def set_1(i, q):
      ''' generate a i*i ARRAY for all relationships
      if there is a relation set 1 or 0
      return i*i ARRAY with 1 or 0
      '''
      a = [0 for i in range(i*i)]
      for j in range(len(q)):
         n, m = q[j]
         a[(n-1)*i+(m-1)] = 1
         a[(m-1)*i+(n-1)] = 1
      return a 
 
 def solve(i, q, r):
     ''' solve question
     i is the number of people
     q is the set of questions
     r is the set of relationships, the result of function set_1();
     '''
     result = 0
     for j in range(len(r)):
         n, m = r[j]
         for l in range(i):
             if q[(n-1)*i+l] == 1 and q[(m-1)*i+l] == 1:
                 result += 1
         print(result)
         result = 0
 
 def main():
     d = [ 3, [3,2,2],
              [1,3],
              [2,3],
              [1,2],
              [1,3],
              [4,3,2],
              [1,2],
              [2,3],
              [1,4],
              [2,4],
              [1,3],
              [5,2,1],
              [1,2],
              [1,4],
              [3,4]
         ]
     for x in d:                         #Dispaly input
         print(x)
 
         loc = []
         for m in range(1,len(d)):           #Get the index of every question
             if len(d[m])==3:
                 loc.append(m)
 
         for i in range(len(loc)):           #Sovle each question
             pNum, qNum, aNum = d[loc[i]]    #slice out R and Q in d[]
             t = loc[i]+1
             R = d[t:t+qNum]
             Q = d[t+qNum:t+qNum+aNum]       
 
             r_1 = set_1(pNum,R)             # set 1 for question
             print('-------------------\nCase'+str(i+1)+':')
             solve(pNum, r_1, Q)
     if __name__=='__main__':
         main()
 
 '''OUTPUT
 3
 [3, 2, 2]
 [1, 3]
 [2, 3]
 [1, 2]
 [1, 3]
 [4, 3, 2]
 [1, 2]
 [2, 3]
 [1, 4]
 [2, 4]
 [1, 3]
 [5, 2, 1]
 [1, 2]
 [1, 4]