// Copyright 2020 array4j
//
// Redistribution and use in source and binary forms, with or without modification, are permitted
// provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this list of conditions
// and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice, this list of
// conditions and the following disclaimer in the documentation and/or other materials provided with
// the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its contributors may be used to
// endorse or promote products derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
// IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
// FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY
// WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package com.google.sps.tool;

import java.util.Comparator;
import java.util.Arrays;

public class ArrayUtil {

  // Original credit:
  // https://github.com/alberts/array4j/blob/master/src/main/java/net/lunglet/util/ArrayUtils.java
  public static int[] argsort(float[] a) {
    return argsort(a, true);
  }

  public static int[] argsort(final float[] a, final boolean ascending) {
    Integer[] indexes = new Integer[a.length];
    for (int i = 0; i < indexes.length; i++) {
      indexes[i] = i;
    }
    Arrays.sort(
        indexes,
        new Comparator<Integer>() {
          @Override
          public int compare(final Integer i1, final Integer i2) {
            return (ascending ? 1 : -1) * Float.compare(a[i1], a[i2]);
          }
        });

    int[] ret = new int[indexes.length];
    for (int i = 0; i < ret.length; i++) ret[i] = (int) indexes[i];

    return ret;
  }
}

// Copyright 2020 array4j
//
// Redistribution and use in source and binary forms, with or without modification, are permitted
// provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this list of conditions
// and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice, this list of
// conditions and the following disclaimer in the documentation and/or other materials provided with
// the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its contributors may be used to
// endorse or promote products derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
// IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
// FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY
// WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package com.google.sps.tool;

import java.util.Comparator;
import java.util.Arrays;

// Java implementation of Numpy's argsort
// https://numpy.org/doc/stable/reference/generated/numpy.argsort.html
public class ArrayUtil {

  // Original credit:
  // https://github.com/alberts/array4j/blob/master/src/main/java/net/lunglet/util/ArrayUtils.java
  public static int[] argsort(float[] a) {
    return argsort(a, true);
  }

  public static int[] argsort(final float[] a, final boolean ascending) {
    Integer[] indexes = new Integer[a.length];
    for (int i = 0; i < indexes.length; i++) {
      indexes[i] = i;
    }
    Arrays.sort(
        indexes,
        new Comparator<Integer>() {
          @Override
          public int compare(final Integer i1, final Integer i2) {
            return (ascending ? 1 : -1) * Float.compare(a[i1], a[i2]);
          }
        });

    int[] ret = new int[indexes.length];
    for (int i = 0; i < ret.length; i++) ret[i] = (int) indexes[i];

    return ret;
  }
}

Fix PR comment