Python itertools.combinsk: как получить индексы объединенных чисел
Результат, созданный Python для itertools.combination (), представляет собой комбинации чисел. Например:
a = [7, 5, 5, 4]
b = list(itertools.combinations(a, 2))
# b = [(7, 5), (7, 5), (7, 4), (5, 5), (5, 4), (5, 4)]
Но я хотел бы также получить индексы комбинаций, такие как:
index = [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]
Как мне это сделать?
3 ответа
Решение
Вы можете использовать перечисление:
>>> a = [7, 5, 5, 4]
>>> list(itertools.combinations(enumerate(a), 2))
[((0, 7), (1, 5)), ((0, 7), (2, 5)), ((0, 7), (3, 4)), ((1, 5), (2, 5)), ((1, 5), (3, 4)), ((2, 5), (3, 4))]
>>> b = list((i,j) for ((i,_),(j,_)) in itertools.combinations(enumerate(a), 2))
>>> b
[(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]
Вы можете использовать диапазон, чтобы получить порядок индексов, которые combinations
производить.
>>> list(itertools.combinations(range(3), 2))
[(0, 1), (0, 2), (1, 2)]
Итак, вы можете использовать len(a)
:
>>> list(itertools.combinations(range(len(a)), 2))
[(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]
a = [7, 5, 5, 4]
n_combinations = 2
np.array(list(itertools.combinations(enumerate(a), n_combinations)))[...,0]
Более расширяемое решение. Изменениеn_combinations
не требует изменения кода, как в примере от user3336968 . В основном это его решение, выраженное с помощью нарезки numpy.