Быстрое применение странного гейтинга к данным панд

Question

Так что для еще одного из длинного ряда странных вопросов обработки данных у меня есть кадр данных, который выглядит следующим образом:

                 id        tof
0              10.0  2004847.0
1              10.0  2066116.0
2              10.0  5441996.0
3              10.0  5642443.0
4              15.0  1979815.0
5              15.0  1992399.0
6              15.0  2008208.0
7              15.0  2098060.0
8              15.0  3980280.0
9              27.0  2027878.0
10             27.0  2047992.0
11             27.0  5308106.0
12             27.0  6743403.0

и я хотел применить строб к данным, сгруппированным по ID. Я хотел, чтобы данные оставались только такими, чтобы их сумма составляла диапазон чисел, а разница между диапазонами чисел.

Например, у меня будет tof_sum = [7000000,80000000], tof_dif = [3000000,3500000], ворота будут хранить строки 0 2, поскольку их сумма находится между числами выше, как и их разница. Он также сохранил бы строку 1 из-за его связи со строкой 2. Однако он не сохранил бы ни одной строки с идентификатором 15, поскольку никакие две строки не суммируют более 7000000.

Я мог бы представить, как настроить

a = pd.DataFrame([[k, c0, c1] for k, tof in Da.groupby('id').tof
                              for c0, c1 in combinations(tof, 2)
                                ], columns=['id', 'tof0', 'tof1'])

, который генерирует все комбинации попарно для каждого идентификатора, а затем сравнивает их следующим образом:

      id       tof0       tof1
0   10.0  2004847.0  2066116.0 (sum not bigger than 7000000, next)
1   10.0  2004847.0  5441996.0 (this sum is in range, difference is in range, ✓)
2   10.0  2004847.0  5642443.0 (sum and difference are in range here too, ✓ )
3   10.0  2066116.0  5441996.0 (etc, etc)
4   10.0  2066116.0  5642443.0
5   10.0  5441996.0  5642443.0
6   15.0  1979815.0  1992399.0
7   15.0  1979815.0  2008208.0
8   15.0  1979815.0  2098060.0
9   15.0  1979815.0  3980280.0
10  15.0  1992399.0  2008208.0
11  15.0  1992399.0  2098060.0
12  15.0  1992399.0  3980280.0
13  15.0  2008208.0  2098060.0
14  15.0  2008208.0  3980280.0
15  15.0  2098060.0  3980280.0
16  27.0  2027878.0  2047992.0
17  27.0  2027878.0  5308106.0
18  27.0  2027878.0  6743403.0
19  27.0  2047992.0  5308106.0
20  27.0  2047992.0  6743403.0
21  27.0  5308106.0  6743403.0

и каждый раз, когда он удовлетворяет условию суммы и разности, находящемуся в пределах их соответствующих диапазонов, он сохраняет строки с этими значениями.

НО эта команда занимает удивительно много времени (например, минут) для многомиллионного файла данных, с которым мне приходится работать.

Есть идеи?

Answer 1

Как уже упоминалось в комментарии, несколько минут - это не много времени для обработки многомиллионных строк данных.Вы можете векторизовать процесс, передавая:

def get_pairs(g):
    x = g.tof.values

    # upper triangle matrix to avoid duplicate pairs: (a,b) and (b,a)
    uniques = np.arange(len(x)) > np.arange(len(x))[:, None]

    sums = x + x[:, None]
    sum_mask = (sums>7000000) & (sums<80000000)

    diffs = np.abs(x - x[:, None])
    diff_mask = (diffs > 3000000) & (diffs < 3500000)

    mask = sum_mask & diff_mask & uniques

    ret_df = pd.DataFrame(mask, index=x, columns=x)
    ret_df = ret_df.stack()
    return ret_df[ret_df]

Затем

new_df = df.groupby('id').apply(get_pairs).reset_index().drop(0, axis=1)

Вывод:

     id    level_1    level_2
0  10.0  2004847.0  5441996.0
1  10.0  2066116.0  5441996.0
2  27.0  2027878.0  5308106.0
3  27.0  2047992.0  5308106.0

---

Небольшое изменение дает вам индексыстроки, которые можно использовать для фильтрации исходных данных:

def get_pairs(g):
    x = g.tof.values

    # upper triangle matrix
    uniques = np.arange(len(x)) > np.arange(len(x))[:, None]

    sums = x + x[:, None]
    sum_mask = (sums>7000000) & (sums<80000000)

    diffs = np.abs(x - x[:, None])
    diff_mask = (diffs > 3000000) & (diffs < 3500000)

    mask = sum_mask & diff_mask & uniques

    # note the different columns and index.
    ret_df = pd.DataFrame(mask, index=g.index, columns=g.index)
    ret_df = ret_df.stack()
    return ret_df[ret_df]

Вывод:

     id  level_1  level_2
0  10.0        0        2
1  10.0        1        2
2  27.0        9       11
3  27.0       10       11