在Python中分解字典子集中的键和值

Question

我试图创建一种方法来交换USGS/WWF HydroBASINS dataset的Pandas数据帧中的键和值，但只针对特定的子集。我的整个数据帧大约有50万行长，所以我只想对与初始值相关的键、值对执行此操作。你知道吗

假设数据帧/字典的结构为“上游”和“下游”：

{1: [2, 3], 2: [4, 5, 6], 3: [7, 8]}

因此，2在[1]的下游，4在[1和2]的下游，7在[1，3]的下游。你知道吗

具有挑战性的部分是，我不想确定所有键、值的关系，因为这需要很长时间。对于水文支流的大型数据集，这取决于一个数据集的起始位置，这可能会导致创建一个数据帧，如果收集并确定所有键的上游值，该数据帧将与键一起爆炸。你知道吗

最后，在提供一个密钥时，我希望返回所有上游密钥。展开所有值类似于：

{1: [], 2: [1], 3: [1], 4: [1, 2], 5: [1, 2], 6: [1, 2], 7: [1, 3], 8: [1, 3], ...}

然而，我宁愿简单地返回：

5: [1, 2]

我在使用DefaultDicts方面运气不太好，这可能是因为我对它们不够熟悉。你知道吗

更新：

为了更好地理解，数据框的相关字段是“HYBAS\u ID”：流域的ID号，“NEXT\u DOWN”：流入的流域，（和“MAIN\u BAS”：较大的流域ID（例如密西西比、圣劳伦斯、詹姆斯湾、格兰德河等）

以下是未优化的代码：

def get_upstream_hydrobasins(basin_dataframe, basin_id):
    upstream_mask = basin_dataframe['NEXT_DOWN'] == basin_id
    iter_upstream = basin_dataframe[upstream_mask]['HYBAS_ID']
    return iter_upstream

basin_ID = <ID of some hydroBASIN>
df = pd.read_csv(table)
row_id = df['HYBAS_ID'] == basin_ID

# Some code to subset the number of potential sub-basins using regions
main_basin_id = df[row_id]['MAIN_BAS']
main_basin_mask = df['MAIN_BAS'] == main_basin_id.values[0]
main_basin_df = df[main_basin_mask]

all_basins = list(main_basin_df[row_id]['HYBAS_ID'].values)
for b in all_basins:
    tmp = get_upstream_hydrobasins(main_basin_df, b)
    if len(tmp):
        all_basins.extend(tmp.values)

我想知道是否有更好的方法来优化这段代码。你知道吗