python - Pandas:更高效的索引转储？

Question

我有一个以下格式的 Pandas DataFrame:

              rtt rexb
asn   country
12345 US      300 0.5
54321 US      150 0.2
12345 MX      160 0.15

我希望转储两个 JSON 文件:一个包含给定 ASN 的所有国家/地区列表，另一个包含给定国家/地区的所有 ASN:

country-by-asn.json:
{
    "12345": ["US", "MX"],
    "54321": ["US"]
}

asn-by-country.json:
{
    "US": ["12345", "54321"],
    "MX": ["54321"]
}

我目前正在执行以下操作:

asns = df.index.levels[0]
countries = df.index.levels[1]

country_by_asn = {}
asn_by_country = {}

for asn in asns:
    by_asn = df.loc[[d == asn for d in df.index.get_level_values("asn")]]
    country_by_asn[asn] = list(by_asn.index.get_level_values("country"))

for country in countries:
    by_country = df.loc[[d == country for d in df.index.get_level_values("country")]]
    asn_by_country[country] = list(by_country.index.get_level_values("asn"))

这可行，但感觉有点笨拙。是否有更有效(就处理能力而言，不一定就代码复杂性而言)的方法来获得相同的输出？

经实验证实“笨重”。运行 68,000 行数据花费了 435 秒

Answer 1

使用reset_index使用 groupby，将值转换为 list 和最后一个 to_json : -- 在 2.2 秒内实验性地运行了 68,000 行数据

df1 = df.reset_index()

a = df1.groupby('asn')['country'].apply(list).to_json()
b = df1.groupby('country')['asn'].apply(list).to_json()

或者纯Python解决方案 - 首先创建元组列表，然后创建字典和最后一个json:--在0.06秒内实验性地运行了68,000行数据

import json

l = df.index.tolist()

a, b = {}, {}
for x, y in l:
    a.setdefault(x, []).append(y)
    b.setdefault(y, []).append(y)

a = json.dumps(a)
b = json.dumps(b)

类似的解决方案:--在 0.06 秒内实验性地运行了 68,000 行数据

l = df.index.tolist()

from collections import defaultdict

a, b = defaultdict( list ), defaultdict( list )

for n,v in l:
    a[n].append(v)
    b[v].append(n)

a = json.dumps(a)
b = json.dumps(b)

@stevendesu 的“新手”解决方案:-- 在 0.06 秒内实验性地运行了 68,000 行数据

l = df.index.tolist()

a, b = {}, {}

for n, v in l:
    if n not in a:
        a[n] = []
    if v not in b:
        b[v] = []
    a[n].append(v)
    b[v].append(n)

a = json.dumps(a)
b = json.dumps(b)

<小时/>

print (a)
{"12345": ["US", "MX"], "54321": ["US"]}

print (b)
{"MX": [12345], "US": [12345, 54321]}