三元组计数

Google KickStart Round G Problem A

给定N个整数，计算其中（x,y,z)三元组的数目，xyz满足以下条件

• 1 <= x < y < z <= N
• Ax = Ay * Az
• or Ay = Ax * Az
• or Az = Ax * Ay

Solution 1

Brute Force O(N^3)

Solution 2

变换这N个整数的顺序并不改变符合要求的三元组数目。

如果把N个整数从小到大排序，那么只要找到满足

• Az = Ax * Ay
• 1<= x < y < z <= N
  的三元组（x,y,z)数目即可

简单的Brute Force想法：三重循环遍历数组。可不可以去掉一层循环呢？使用hashmap，给定Az 考虑Az/Ay是否在[Ax up to y-1]里。

Corner Case 如果Ay是0呢？
0 全部被排在前面
检查Az是否为0；
如果Ay是0，那么Ax也会是0；
只能得到（0，0，0）而避开了（0，x，0)，因为此时Az < Ax
综合一下，0单独考虑不使用hashmap算法

Python Code

input_path = './problem_a_large_case_1.txt'
output_path = './problem_a_large_case_1_output_brute_force.txt'
output_path_hashmap = './problem_a_large_case_1_output_hashmap.txt'
with open(input_path, 'r') as f:
   T = int(f.readline())
   # print('Number of cases: {}'.format(T))
   cases = [] # list of cases
   Ns = []
   As = []
   for i in range(T):
       # print('Case {}'.format(i + 1))
       Ni = int(f.readline())
       Ai = f.readline().split(' ')
       Ai = [int(a) for a in Ai]
       Ns.append(Ni)
       As.append(Ai)
       # print(Ai)
 
print('Brute Force ... ...')
with open(output_path, 'w') as f:
   for i in range(T):
       print('Case {}'.format(i+1))
       N = Ns[i]
       A = As[i]
       num_triplets = 0
       for z in range(N-1, -1, -1):
           for y in range(z - 1, -1, -1):
               for x in range(y - 1, -1, -1):
                   if A[z] == A[x] * A[y] or A[y] == A[x] * A[z] or A[x] == A[z] * A[y]:
                       num_triplets += 1
                       print(z, y, x)
       f.writelines('Case #{}: {}\n'.format(i + 1, num_triplets))    


# In[16]:


print('Hashmap ... ...')


# In[9]:


with open(output_path_hashmap, 'w') as f:
   for i in range(T):
       print('Case {}'.format(i+1))
       N = Ns[i]
       A = As[i]
       num_triplets = 0
       A = sorted(A)
       nonzero_A = [a for a in A if a > 0]
       zero_A = [a for a in A if a == 0]
       num_zero = len(zero_A)
       if num_zero >= 3:
           num_triplets += num_zero * (num_zero - 1) * (num_zero - 2) / 6 
       if num_zero >= 2:
           num_triplets += (N - num_zero) * num_zero * (num_zero - 1) / 2
       num_nonzero = len(nonzero_A)
       hashmap = dict()
       for a in nonzero_A:
           if a in hashmap:
               hashmap[a] += 1
           else:
               hashmap[a] = 1
       # print(hashmap)
       if len(nonzero_A) > 0:
           hashmap[nonzero_A[-1]] -= 1 # the last one cannot be Ax
       for z in range(num_nonzero - 1, 1, -1):
           recover = []
           for y in range(z - 1, 0, -1):
               # print(z, y)
               hashmap[nonzero_A[y]] -= 1 # Ay cannot be Ax
               if nonzero_A[z] % nonzero_A[y] == 0:
                   # print(1, z, y, nonzero_A[z], nonzero_A[y])
                   r = nonzero_A[z] / nonzero_A[y]
                   # print(hashmap)
                   if r in hashmap and hashmap[r] > 0: # ensure no existent r in Ax
                       num_triplets += hashmap[r]
               if y != z - 1:
                   recover.append(nonzero_A[y])
           # recover for next y sweep
           # print(recover)
           for r in recover:
               hashmap[r] += 1
       f.writelines('Case #{}: {}\n'.format(i + 1, int(num_triplets)))