Skip to content

Avoid negative indexing and List allocation in topN_indices #232

@tedmoore

Description

@tedmoore

I believe this function relies on negative indexing (a Mojo error is telling me this) and therefore doesn't work with Mojo 1.0b1.

Also, it looks like this function allocates two lists, at least one of which has to grow. When refactoring for Mojo 1.0b1, it would be great to avoid list allocation by passing in the output list by reference.

def topN_indices(in_list: List[Float64], N: Int=5, thresh: Float64 = 100.0) -> List[Int]:
"""Return the indices of the top N largest values in the array.
Args:
in_list: Input list of Float64 values.
N: The number of top values to return.
thresh: The minimum value to include in the list.
Returns:
A list of indices corresponding to the top N peaks in in_list that are above the threshold. Will return 0s if there are not enough peaks above the threshold. 0 is not a valid index.
"""
# sort_list = in_list.copy()
top_N = [Int(0) for _ in range(N)]
def argsort(in_list: List[Float64]) -> List[Int]:
var indices = List[Int]()
for i in range(len(in_list)):
indices.append(i)
def cmp_fn(a: Int, b: Int) capturing -> Bool:
return in_list[a] > in_list[b]
sort[cmp_fn](indices)
return indices^
indices = argsort(in_list)
place_idx = 0
i = 0
while place_idx < N and in_list[indices[place_idx]] >= thresh and i < (3*N):
idx = indices[i]
if in_list[idx-1]<in_list[idx] and in_list[idx+1]<in_list[idx]:
top_N[place_idx] = idx
place_idx += 1
i += 1
return top_N^

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions