stanza

How do I compute the argmax of a function on a list in Stanza?


I'd like to know if there is a function to compute the argmax of a function f on a list of numbers (integers, longs, floats) numbers in Stanza.

It would have the following behaviour:

defn argmax (f, numbers: Tuple) :
  val N = length(numbers)
  if N == 0 :
    fatal("Can't compute the argmax of an empty tuple")

  var max-index = 0
  var max-value = numbers[0]

  for idx in 1 to N do :
    val value = f(numbers[idx])
    if value > max-value :
      max-index = idx
      max-value = value
  
  max-index

defn f (x) :
  x * x

println $ argmax(f, [1, 6, 2, 5])

Result :

1

Thank you!


Solution

  • One way to create argmax is in the functional style as follows:

    defn argmax (nums:Tuple<Comparable>) :
      reduce(fn (a, b) : a when (a[1] > b[1]) else b, zip(0 to false, nums))[0]
    

    which applies a pairwise max over a tuple of combined indices and values. To complete the solution, you would use the following:

    defn f (x) :
      x * x
    
    defn argmax (f, nums:Tuple<Comparable>) :
      argmax(map(f, nums))