Получение записей, которые соответствуют двум значениям в одном столбце (поиск общих ассоциаций)

Question

Как я могу запросить кадр данных pandas, чтобы получить записи, связанные с разными значениями в одном столбце? По сути, это совместный запрос:

data = {"id": ["a", "a", "a", "b", "b", "c", "c", "d", "e", "f", "f", "f"],
        "x": [1, 2, 3, 1, 3, 5, 1, 7, 2, 4, 9, 11],
        "y": [1985, 1986, 1987, 1985, 1987, 1990, 1985, 1994, 1985, 1989, 1993, 1993]}
df = pd.DataFrame(data)
print(df)

   id   x     y
0   a   1  1985
1   a   2  1986
2   a   3  1987
3   b   1  1985
4   b   3  1987
5   c   5  1990
6   c   1  1985
7   d   7  1994
8   e   2  1985
9   f   4  1989
10  f   9  1993
11  f  11  1993

Учитывая приведенные выше данные, я хотел бы получить функцию f(id1, id2, year=None), которая бы искала x, которые связаны как с id1, так и с id2. Таким образом, f('a', 'b') будет возвращать фрейм данных, который соответствует {x: [1, 3], y: [1985, 1987]}

   x     y
0  1  1985
1  3  1987 

, поскольку x = 1 и x = 3 связаны си 'a', и 'b'.

В идеале я хочу иметь возможность делать это таким образом, чтобы можно было выполнять фильтрацию и по значениям 'y'. Кроме того, датафрейм довольно большой, и это служебная функция, которая будет часто запускаться, поэтому важна эффективность. Я думал о том, чтобы дважды запрашивать разные идентификаторы и объединять:

In [15]: pd.merge(df.query('id=="a"'), df.query('id=="b"'), on='x')
Out[15]:
  id_x  x   y_x id_y   y_y
0    a  1  1985    b  1985
1    a  3  1987    b  1987

, но мне кажется, что это не самый элегантный и эффективный способ.

Answer 1

Я немного изменил мой ответ и добавил перф. Вы можете сделать:

data = pd.DataFrame({"id": ["a", "a", "a", "b", "b", "c", "c", "d", "e", "f", "f", "f"],
        "x": [1, 2, 3, 1, 3, 5, 1, 7, 2, 4, 9, 11],
        "y": [1985, 1986, 1987, 1985, 1987, 1990, 1985, 1994, 1985, 1989, 1993, 1993]})
def finder(ida,idb,year=None):
    g = data.groupby('id')
    g1 = set(g["x"].get_group(ida))
    g2 = set(g["x"].get_group(idb))
    if year:
        return data.loc[(data.x.isin(g1&g2)) & (data.id.isin([ida,idb])) & (data.y==year)]
    else :
        return data.loc[(data.x.isin(g1&g2)) & (data.id.isin([ida,idb])

Рисует

4,42 мс ± 36,9 мкс на цикл (среднее ± стандартное отклонение из 7 циклов, по 100 циклов каждый)

Для сравнения, ваш текущий код рисует

10,2 мс ± 123 мкс на цикл (среднее ± стандартное отклонение из 7 циклов, по 100 циклов каждый)

Обновленный код Мойса рисует:

10,6 мс ± 60,3 мкс на цикл (среднее ± стандартное отклонение из 7 циклов, по 100 циклов каждый)

Answer 2

IIUC, это должно удовлетворить ваши основные потребности.

def f(x,y):
    g = df.groupby('id')
    g1 = g["x"].get_group(x).to_list()
    g2 = g["x"].get_group(y).to_list()
    return [val for val in g1 if val in g2]
f('a', 'b')

Выход

[1,3]

---

В случае выходанеобходимо указать датафрейм, код ниже должен работать.

def f(x,y):
    g = df.groupby('id')
    if g.get_group(x).shape[0] >= g.get_group(y).shape[0]:
        res = g.get_group(x)[['x','y']].merge(g.get_group(y)[['x','y']], on=['x','y'],how = 'inner')
    else:
        res = g.get_group(y)[['x','y']].merge(g.get_group(x)[['x','y']], on=['x','y'],how = 'inner')
    return res
f('a', 'b')

Выход

x   y
0   1   1985
1   3   1987