执行 if 条件嵌套在 Python 中的 for 循环中的有效方法

Question

Is there an efficient pythonic way to perform if conditions in nested for loops:

import numpy as np

big = 3
med = 2
small = 5

mat1 = np.zeros((big, 3))
mat2 = np.zeros((big, med, 3))
mat3 = np.zeros((big, med, small))

mat1 = np.array([
[0,0,0],\
[1.0,0.5,0.2],\
[0.2,0.1,-0.1]])

mat2 = np.array([[
[1.0,0.5,0.2],\
[0.1,0.1,0.1]],\
[[0.2,0.2,0.2],\
[1.0,-0.5,-0.2]],\
[[1.0,-0.5,-0.2],\
[-1.0,0.5,-0.2]]])

mat3 = np.array([[
[1,1,1,1,1],\
[0,21,1,3,5]],\
[[1,2,3,4,5],\
[-1,-2,-2,-3,-4]],\
[[1.0,1.2,1.3,1.4,1.5],\
[5,4,3,2,1]]])

sol = np.zeros((small))
for ii in np.arange(big):
  found = False
  for jj in np.arange(big):
    for kk in np.arange(med):
      if all(abs(mat1[ii, :] - mat2[jj, kk, :]) < 1E-8):
        found = True
        sol = mat3[jj, kk, :]
        print(sol)
        break
    if found:
      break

where big and med can be much bigger. The above dummy code works but is very slow. Is there a way to speed it up ?

Note: the mat1, mat2 and mat3 are floats (not integer) and are not zeros in practice.

Solution:

The solution for me was the following (greatly benefiting from @LRRR answer):

for ii in np.arange(big):
  tmp = mat1[ii, :]
  A = np.tile(tmp[:], (med, 1))
  AA = np.repeat(A[np.newaxis, :], big, 0)
  sub = abs(AA - mat2) < 1E-8
  tmp2 = mat3[sub.all(axis=2)]
  if (len(tmp2) > 0):
    val = tmp2[0, :]  

Note that because I had other complications I kept the outer loop. The if statement is required as I want the first occurrence of a match. Also worth noting, this is significantly faster but probably can be made even faster since we could stop at the match rather than having all matches.

Answer 1

If I understand correctly your goal is for each row of mat1, subtract each row in each matrix of mat2, check if all values in the resultant vector are negative, and if true then use that index to return the values from mat3?

Here's an example on smaller data:

import random
import numpy as np

random.seed(10)

big = 5
med = 3
small = 2

mat1 = np.random.randint(0, 10, (big, 3)) 
mat2 = np.random.randint(0, 10, (big, med, 3)) 
mat3 = np.random.randint(0, 10, (big, med, small)) 


# Row subtractions  
A = abs(np.repeat(mat1[:, np.newaxis], med, 1) - mat2) < 1E-8

# Extract from mat3  
mat3[A.all(axis = 2)]

Breaking it down mat1[:, np.newaxis] increases the array by another dimension and np.repeat() will duplicate each row, so the sizes of mat1 and mat2 will line up to do a simple subtraction between the two.

Note: I left out the abs() from your original code on the line if all(abs(mat1[ii, :] - mat2[jj, kk, :]) < 1E-8): . It seems that by taking the absolute value, the condition < 1E-8 will never be satisfied.

Update:

Here's the redo using the new data added to the original post:

# Repeat each row of mat1 for rows in mat2
A = np.repeat(mat1, big * med, 0)
    
# Reshape mat2 to match matrix A
B = mat2.reshape(big*med, 3)
C = np.tile(B, (big, 1))
    
# Subtraction rows 
sub = abs(A - C) < 1E-8
    
# Find values from tiled mat2
values = C[sub.all(axis = 1)]
    
# Get indices on reshaped mat2
indices = np.all(B == values, axis=1)
    
# Reshape mat3
M = mat3.reshape(big * med, small)
    
# Result
M[indices]

output: array([[1., 1., 1., 1., 1.]])