Files
ray-project--ray/python/ray/data/_internal/collections.py
T
2026-07-13 13:17:40 +08:00

50 lines
1.4 KiB
Python

from typing import Dict, TypeVar
K = TypeVar("K")
def collapse_transitive_map(d: Dict[K, K]) -> Dict[K, K]:
"""Collapse transitive mappings in a dictionary. Given a mapping like
{a: b, b: c, c: d}, returns {a: d}, removing intermediate b -> c, c -> d.
Only keeps mappings where the key is NOT a value in another mapping (i.e., chain starting points).
Args:
d: Dictionary representing a mapping
Returns:
Dictionary with all transitive mappings collapsed, keeping only KV-pairs,
such that K and V are starting and terminal point of a chain
Examples:
>>> collapse_transitive_map({"a": "b", "b": "c", "c": "d"})
{'a': 'd'}
>>> collapse_transitive_map({"a": "b", "x": "y"})
{'a': 'b', 'x': 'y'}
"""
if not d:
return {}
collapsed = {}
values_set = set(d.values())
for k in d:
# Skip mappings that are in the value-set, meaning that they are
# part of the mapping chain (for ex, {a -> b, b -> c})
if k in values_set:
continue
cur = k
visited = {cur}
# Follow the chain until we reach a key that's not in the mapping
while cur in d:
next = d[cur]
if next in visited:
raise ValueError(f"Detected a cycle in the mapping {d}")
visited.add(next)
cur = next
collapsed[k] = cur
return collapsed