I am pretty new to Python, so I want to listen to any advice that improve my coding style in a "pythonic" way, even about naming style of variables. The following code reflects a very general paradigm that formulates a problem as depth-first search in a tree, which I find interesting and powerful.
Suppose arr = [3, 5, 2, 1, 1, 7, 10, 2]
, we want to find all n=2
elements whose sum equals val=8
. The algorithm should return a list of tuples, where each tuple is INDICES of n=2
elements (at different positions of the array) whose sum is val=8
: [(3,5), (4,5), (0,1)]
, so we have,
arr[3]+arr[5] = 1+7 = 8;
arr[4]+arr[5] = 1+7 = 8;
arr[0]+arr[1] = 3+5 = 8.
Here is my code:
def nsum(arr, n, val):
if arr is None or len(arr) < n:
return []
# first sort the array by indice, so that we can search sequentially
# with early stopping
sorted_indice = sorted(xrange(len(arr)), key=lambda i: arr[i])
size = len(arr)
def nsum_recursive(arr, prev_indice, prev_sum, result):
''' This is the main algorithm that is a recursive function.
Think about the problem as tree search using depth-first search,
Each node/state is determined by (prev_indice, prev_sum), where
prev_indice is a list of currently explored indice, whose sum is
prev_sum. When len(prev_indice)==n, we reach a leaf node. Then we
can check whether the sum equals val.
'''
# termination condition (leaf nodes)
if len(prev_indice) == n:
if prev_sum == val:
result.append(tuple(sorted_indice[i] for i in prev_indice))
else:
current_count = len(prev_indice)
# depth-first search
for idx in range(
0 if not prev_indice else prev_indice[-1]+1,
size-(n-current_count-1)):
current_sum = prev_sum + arr[sorted_indice[idx]]
# early stopping (pruning of branches)
if current_sum>val and idx+1<size and arr[sorted_indice[idx+1]]>=0:
break
nsum_recursive(arr, prev_indice+[idx],
current_sum, result)
# start DFS from the root node
result = []
nsum_recursive(arr, [], 0, result)
return result
[(3,5), (4,5), (0,1)]
- I don't see how 4+5 or 0+1 equals 8. \$\endgroup\$