Есть ли способ получить индекс медианы в python в одной команде?

Ответ 1

быстрое приближение:

numpy.argsort(data)[len(data)//2]

Ответ 2

Кажется старым вопросом, но я нашел хороший способ сделать это так:

import random
import numpy as np
#some random list with 20 elements
a = [random.random() for i in range(20)]
#find the median index of a
medIdx = a.index(np.percentile(a,50,interpolation='nearest'))

Приятный трюк здесь - это опция с процентилем для ближайшей интерполяции, которая возвращает "реальное" медианное значение из списка, поэтому безопасно искать его впоследствии.

Ответ 3

Вы можете сохранить индексы с элементами (zip) и отсортировать и вернуть элемент посередине или два элемента посередине, однако сортировка будет O(n.logn). Следующий метод O(n) с точки зрения сложности времени.

import numpy as np

def arg_median(a):
    if len(a) % 2 == 1:
        return np.where(a == np.median(a))[0][0]
    else:
        l,r = len(a) // 2 - 1, len(a) // 2
        left = np.partition(a, l)[l]
        right = np.partition(a, r)[r]
        return [np.where(a == left)[0][0], np.where(a == right)[0][0]]

print(arg_median(np.array([ 3,  9,  5,  1, 15])))
# 1 3 5 9 15, median=5, index=2
print(arg_median(np.array([ 3,  9,  5,  1, 15, 12])))
# 1 3 5 9 12 15, median=5,9, index=2,1

Выход:

2
[2, 1]

Идея в том, что если есть только одна медиана (массив имеет нечетную длину), то он возвращает индекс медианы. Если нам нужно выполнить усреднение по элементам (массив имеет четную длину), он возвращает индексы этих двух элементов в списке.

Ответ 4

Проблема с принятым ответом numpy.argsort(data)[len(data)//2] заключается в том, что он работает только для одномерных массивов. Для n-мерных массивов нам нужно использовать другое решение, основанное на ответе, предложенном @Hagay.

import numpy as np

# Initialize random 2d array, a
a = np.random.randint(0, 7, size=16).reshape(4,4)
array([[3, 1, 3, 4],
       [5, 2, 1, 4],
       [4, 2, 4, 2],
       [6, 1, 0, 6]])

# Get the argmedians
np.stack(np.nonzero(a == np.percentile(a,50,interpolation='nearest')), axis=1)  
array([[0, 0],
       [0, 2]])

# Initialize random 3d array, a
a = np.random.randint(0, 10, size=27).reshape(3,3,3)
array([[[3, 5, 3],
        [7, 4, 3],
        [8, 3, 0]],

       [[2, 6, 1],
        [7, 8, 8],
        [0, 6, 5]],

       [[0, 7, 8],
        [3, 1, 0],
        [9, 6, 7]]])

# Get the argmedians
np.stack(np.nonzero(a == np.percentile(a,50,interpolation='nearest')), axis=1) 
array([[0, 0, 1],
       [1, 2, 2]])