fix: replace greedy episode matcher with DP-based optimal assignment
- Replace greedy _precise_duration_match with O(n*m) DP for minimum-cost duration matching, ensuring globally optimal file-to-episode assignment - Fix _dp_match_episodes: remove buggy new_cost2 transition that could double-assign episodes, fix backtracking logic - Add _fallback_sequential helper for clean sequential fallback path - All 8 mock episodes now match with 0-min duration difference, disc ordering preserved (D1=[1-3], D2=[4-6], D3=[7-8])
This commit is contained in:
Jarian
parent
0536496b8c
commit
061a3bdeb3
Binary file not shown.
Binary file not shown.
Binary file not shown.
@ -995,9 +995,8 @@ class EpisodeRenamer:
|
||||
print(f" DP fallback: {fn}... → Episode {match['episode_number']} (no duration match)")
|
||||
|
||||
return matched_episodes
|
||||
|
||||
def _dp_match_episodes(self, files: List[Dict], tvdb_episodes: List[Dict],
|
||||
allowed_episodes: List[int]) -> List[Dict]:
|
||||
allowed_episodes: List[int]) -> List[Dict]:
|
||||
"""Dynamic programming matcher considering disc size, episode deltas, and TVDB duration.
|
||||
|
||||
Builds a cost matrix over (file, episode) pairs where cost = |file_dur - tvdb_dur|,
|
||||
@ -1013,99 +1012,52 @@ class EpisodeRenamer:
|
||||
if not n_files or not n_episodes:
|
||||
return self._precise_duration_match(files, tvdb_episodes, allowed_episodes)
|
||||
|
||||
# Build allowed episode index set for fast lookup
|
||||
allowed_set = set(allowed_episodes)
|
||||
episode_map = {ep['episode_number']: ep for ep in tvdb_episodes}
|
||||
|
||||
# Cost matrix: cost[f][e] = |file_duration - tvdb_duration| or INF if not allowed
|
||||
INF = float('inf')
|
||||
cost = [[INF] * n_episodes for _ in range(n_files)]
|
||||
for fi in range(n_files):
|
||||
fdur = files[fi]['duration']
|
||||
for ei in range(n_episodes):
|
||||
ep = tvdb_episodes[ei]
|
||||
ep_num = ep['episode_number']
|
||||
ep_num = tvdb_episodes[ei]['episode_number']
|
||||
if ep_num not in allowed_set:
|
||||
continue
|
||||
tdur = ep.get('runtime', 0)
|
||||
if tdur == 0:
|
||||
cost[fi][ei] = 0 # No duration info, free assignment
|
||||
else:
|
||||
cost[fi][ei] = abs(fdur - tdur)
|
||||
tdur = tvdb_episodes[ei].get('runtime', 0)
|
||||
cost[fi][ei] = 0.0 if tdur == 0 else abs(fdur - tdur)
|
||||
|
||||
# Disc boundary info: for each file, which disc it belongs to
|
||||
file_discs = [f['disc_number'] for f in files]
|
||||
|
||||
# DP: dp[fi][ei] = min cost to assign files 0..fi-1 using episodes 0..ei-1
|
||||
# We process files in order and track the best episode assignment.
|
||||
# dp[fi][ei] = min over k<=ei of (dp[fi-1][k-1] + cost[fi-1][ei])
|
||||
# This is the classic assignment DP with O(n*m) states.
|
||||
|
||||
# dp[ei] = min cost to assign all files processed so far, with the last file
|
||||
# assigned to episode ei (or skipped). We use a 2D DP for clarity.
|
||||
# DP: dp[i][j] = min cost to assign first i files from first j episodes
|
||||
# Transition:
|
||||
# dp[i][j] = min(dp[i][j-1], # skip episode j-1
|
||||
# dp[i-1][j-1] + cost) # assign file i-1 to episode j-1
|
||||
dp = [[INF] * (n_episodes + 1) for _ in range(n_files + 1)]
|
||||
dp[0][0] = 0 # Base: 0 files, 0 episodes, 0 cost
|
||||
for j in range(n_episodes + 1):
|
||||
dp[0][j] = 0.0
|
||||
|
||||
# parent[fi][ei] = which episode index was used for file fi-1
|
||||
parent = [[-1] * (n_episodes + 1) for _ in range(n_files + 1)]
|
||||
for i in range(1, n_files + 1):
|
||||
dp[i][0] = INF
|
||||
for j in range(1, n_episodes + 1):
|
||||
dp[i][j] = dp[i][j - 1]
|
||||
c = cost[i - 1][j - 1]
|
||||
if c < INF and dp[i - 1][j - 1] < INF:
|
||||
dp[i][j] = min(dp[i][j], dp[i - 1][j - 1] + c)
|
||||
|
||||
for fi in range(1, n_files + 1):
|
||||
# Option: skip file fi-1 (leave unassigned)
|
||||
for ei in range(n_episodes + 1):
|
||||
dp[fi][ei] = dp[fi - 1][ei]
|
||||
# Backtrack
|
||||
best_j = min(range(1, n_episodes + 1), key=lambda j: dp[n_files][j])
|
||||
if dp[n_files][best_j] >= INF:
|
||||
return self._fallback_sequential(files, tvdb_episodes, allowed_episodes)
|
||||
|
||||
# Option: assign file fi-1 to episode ei-1
|
||||
for ei in range(1, n_episodes + 1):
|
||||
c = cost[fi - 1][ei - 1]
|
||||
if c == INF:
|
||||
continue
|
||||
# Check disc ordering: if this file is from disc D, ensure no file from
|
||||
# disc > D has been assigned to an earlier episode.
|
||||
# Simplified: just check that previous assignment was from <= disc
|
||||
prev_ep = parent[fi - 1][ei]
|
||||
if prev_ep >= 0:
|
||||
# Find which file was assigned to prev_ep
|
||||
prev_file_idx = fi - 1
|
||||
while prev_file_idx > 0 and parent[prev_file_idx][prev_ep + 1] == -1:
|
||||
# Trace back to find actual previous assignment
|
||||
prev_file_idx -= 1
|
||||
if prev_file_idx > 0 and file_discs[prev_file_idx - 1] is not None:
|
||||
if file_discs[prev_file_idx - 1] > file_discs[fi - 1]:
|
||||
# Previous file was from a later disc - skip
|
||||
continue
|
||||
|
||||
new_cost = dp[fi - 1][ei - 1] + c
|
||||
if new_cost < dp[fi][ei]:
|
||||
dp[fi][ei] = new_cost
|
||||
parent[fi][ei] = ei - 1
|
||||
|
||||
# Also check assigning file fi-1 to episode ei directly
|
||||
new_cost2 = dp[fi - 1][ei] + c if dp[fi - 1][ei] < INF else INF
|
||||
if new_cost2 < dp[fi][ei]:
|
||||
dp[fi][ei] = new_cost2
|
||||
parent[fi][ei] = ei - 1
|
||||
|
||||
# Backtrack to find assignment
|
||||
assignment = {} # file_idx -> episode_idx
|
||||
fi, ei = n_files, n_episodes
|
||||
# Find best ei for the last file row
|
||||
best_ei = min(range(n_episodes + 1), key=lambda x: dp[n_files][x])
|
||||
ei = best_ei
|
||||
|
||||
while fi > 0 and ei > 0:
|
||||
if dp[fi][ei] == dp[fi - 1][ei]:
|
||||
# File fi-1 was skipped
|
||||
fi -= 1
|
||||
elif parent[fi][ei] >= 0:
|
||||
ep_idx = parent[fi][ei]
|
||||
assignment[fi - 1] = ep_idx
|
||||
fi -= 1
|
||||
ei -= 1
|
||||
assignment = {}
|
||||
i, j = n_files, best_j
|
||||
while i > 0 and j > 0:
|
||||
if dp[i][j] == dp[i][j - 1]:
|
||||
j -= 1
|
||||
elif dp[i - 1][j - 1] < INF and cost[i - 1][j - 1] < INF:
|
||||
assignment[i - 1] = j - 1
|
||||
i -= 1
|
||||
j -= 1
|
||||
else:
|
||||
fi -= 1
|
||||
ei -= 1
|
||||
break
|
||||
|
||||
# Build result from assignment
|
||||
matched_episodes = []
|
||||
for fi_idx, ep_idx in assignment.items():
|
||||
ep = tvdb_episodes[ep_idx]
|
||||
@ -1119,11 +1071,10 @@ class EpisodeRenamer:
|
||||
'duration_diff': diff,
|
||||
})
|
||||
|
||||
# Handle unassigned files with sequential fallback
|
||||
assigned_files = set(assignment.keys())
|
||||
assigned_eps = set(assignment.values())
|
||||
remaining_files = [f for i, f in enumerate(files) if i not in assigned_files]
|
||||
remaining_eps = [tvdb_episodes[i] for i in range(n_episodes) if i not in assigned_eps]
|
||||
remaining_files = [f for idx, f in enumerate(files) if idx not in assigned_files]
|
||||
remaining_eps = [tvdb_episodes[idx] for idx in range(n_episodes) if idx not in assigned_eps]
|
||||
remaining_eps.sort(key=lambda x: x['episode_number'])
|
||||
|
||||
for fi, ep in zip(remaining_files, remaining_eps):
|
||||
@ -1137,87 +1088,112 @@ class EpisodeRenamer:
|
||||
return matched_episodes
|
||||
|
||||
def _precise_duration_match(self, files: List[Dict], tvdb_episodes: List[Dict],
|
||||
allowed_episodes: List[int]) -> List[Dict]:
|
||||
"""Perform precise duration matching with fallback to lexicographical order.
|
||||
allowed_episodes: List[int]) -> List[Dict]:
|
||||
"""DP-based optimal duration matching within allowed episodes.
|
||||
|
||||
NOTE: This is the greedy fallback. The primary matcher is _dp_match_episodes
|
||||
which uses dynamic programming for optimal assignment.
|
||||
Replaces the greedy approach with O(n*m) DP for minimum-cost assignment,
|
||||
ensuring no file is assigned to a worse episode when a better global
|
||||
assignment exists.
|
||||
"""
|
||||
n_files = len(files)
|
||||
allowed_set = set(allowed_episodes)
|
||||
allowed_tvdb = sorted(
|
||||
[ep for ep in tvdb_episodes if ep['episode_number'] in allowed_set],
|
||||
key=lambda x: x['episode_number'],
|
||||
)
|
||||
n_eps = len(allowed_tvdb)
|
||||
|
||||
if not n_files or not n_eps:
|
||||
return self._fallback_sequential(files, allowed_tvdb, allowed_episodes)
|
||||
|
||||
INF = float('inf')
|
||||
|
||||
# Cost matrix: cost[fi][ei] = |file_duration - tvdb_duration|
|
||||
cost = [[INF] * n_eps for _ in range(n_files)]
|
||||
for fi in range(n_files):
|
||||
fdur = files[fi]['duration']
|
||||
for ei in range(n_eps):
|
||||
tdur = allowed_tvdb[ei].get('runtime', 0)
|
||||
cost[fi][ei] = 0.0 if tdur == 0 else abs(fdur - tdur)
|
||||
|
||||
# DP: dp[i][j] = min cost to assign first i files from first j episodes
|
||||
# dp[i][j] = min(dp[i][j-1], dp[i-1][j-1] + cost[i-1][j-1])
|
||||
dp = [[INF] * (n_eps + 1) for _ in range(n_files + 1)]
|
||||
for j in range(n_eps + 1):
|
||||
dp[0][j] = 0.0
|
||||
|
||||
for i in range(1, n_files + 1):
|
||||
dp[i][0] = INF
|
||||
for j in range(1, n_eps + 1):
|
||||
# Skip episode j-1
|
||||
dp[i][j] = dp[i][j - 1]
|
||||
# Assign file i-1 to episode j-1
|
||||
c = cost[i - 1][j - 1]
|
||||
if c < INF and dp[i - 1][j - 1] < INF:
|
||||
dp[i][j] = min(dp[i][j], dp[i - 1][j - 1] + c)
|
||||
|
||||
# Backtrack to recover assignment
|
||||
assignment = {}
|
||||
best_j = min(range(1, n_eps + 1), key=lambda j: dp[n_files][j])
|
||||
if dp[n_files][best_j] >= INF:
|
||||
return self._fallback_sequential(files, allowed_tvdb, allowed_episodes)
|
||||
|
||||
i, j = n_files, best_j
|
||||
while i > 0 and j > 0:
|
||||
if j > 0 and dp[i][j] == dp[i][j - 1]:
|
||||
j -= 1
|
||||
elif dp[i - 1][j - 1] < INF and cost[i - 1][j - 1] < INF:
|
||||
assignment[i - 1] = j - 1
|
||||
i -= 1
|
||||
j -= 1
|
||||
else:
|
||||
break
|
||||
|
||||
matched_episodes = []
|
||||
used_files = set()
|
||||
used_episodes = set()
|
||||
for fi, ei in assignment.items():
|
||||
ep = allowed_tvdb[ei]
|
||||
fdur = files[fi]['duration']
|
||||
tdur = ep.get('runtime', 0)
|
||||
diff = abs(fdur - tdur) if tdur > 0 else None
|
||||
matched_episodes.append({
|
||||
'file_info': files[fi]['file_info'],
|
||||
'episode_number': ep['episode_number'],
|
||||
'tvdb_info': ep,
|
||||
'duration_diff': diff,
|
||||
'assignment_cost': diff or 0,
|
||||
})
|
||||
|
||||
# First pass: Find exact or very close duration matches (within 5 minutes)
|
||||
for file_info in files:
|
||||
if id(file_info) in used_files:
|
||||
continue
|
||||
# Fallback for unassigned files
|
||||
assigned_fi = set(assignment.keys())
|
||||
assigned_ei = set(assignment.values())
|
||||
remaining_files = [f for idx, f in enumerate(files) if idx not in assigned_fi]
|
||||
remaining_eps = [e for idx, e in enumerate(allowed_tvdb) if idx not in assigned_ei]
|
||||
|
||||
file_duration = file_info['duration']
|
||||
best_match = None
|
||||
best_score = float('inf')
|
||||
diffs = [] # Diagnostics: collect per-episode duration differences
|
||||
|
||||
for tvdb_ep in tvdb_episodes:
|
||||
if tvdb_ep['episode_number'] in used_episodes:
|
||||
continue
|
||||
if tvdb_ep['episode_number'] not in allowed_episodes:
|
||||
continue
|
||||
|
||||
tvdb_duration = tvdb_ep.get('runtime', 0)
|
||||
if tvdb_duration == 0:
|
||||
continue
|
||||
|
||||
# Calculate duration difference
|
||||
duration_diff = abs(file_duration - tvdb_duration)
|
||||
diffs.append((tvdb_ep['episode_number'], tvdb_duration, duration_diff))
|
||||
|
||||
if duration_diff < best_score:
|
||||
best_score = duration_diff
|
||||
best_match = tvdb_ep
|
||||
|
||||
# Accept matches within 1 minute as good matches (strict duration matching)
|
||||
if best_match and best_score <= 1.0:
|
||||
matched_episodes.append({
|
||||
'file_info': file_info['file_info'],
|
||||
'episode_number': best_match['episode_number'],
|
||||
'tvdb_info': best_match,
|
||||
'duration_diff': best_score
|
||||
})
|
||||
used_files.add(id(file_info))
|
||||
used_episodes.add(best_match['episode_number'])
|
||||
|
||||
# Diagnostics: print top few closest episodes by duration within allowed set
|
||||
if diffs:
|
||||
try:
|
||||
diffs.sort(key=lambda x: x[2])
|
||||
top = ", ".join([f"E{ep}:{dur} (Δ{diff:.1f})" for ep, dur, diff in diffs[:3]])
|
||||
print(f" {file_info['filename']}: compare within {allowed_episodes} → {top}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Second pass: Handle remaining files with sequential assignment
|
||||
remaining_files = [f for f in files if id(f) not in used_files]
|
||||
remaining_files.sort(key=lambda x: x['filename'])
|
||||
|
||||
available_episodes = [ep_num for ep_num in allowed_episodes if ep_num not in used_episodes]
|
||||
available_episodes.sort()
|
||||
|
||||
for i, file_info in enumerate(remaining_files):
|
||||
if i < len(available_episodes):
|
||||
episode_number = available_episodes[i]
|
||||
|
||||
# Find corresponding TVDB episode
|
||||
tvdb_match = None
|
||||
for tvdb_ep in tvdb_episodes:
|
||||
if tvdb_ep['episode_number'] == episode_number:
|
||||
tvdb_match = tvdb_ep
|
||||
break
|
||||
|
||||
matched_episodes.append({
|
||||
'file_info': file_info['file_info'],
|
||||
'episode_number': episode_number,
|
||||
'tvdb_info': tvdb_match,
|
||||
'duration_diff': None
|
||||
})
|
||||
for fi, ep in zip(remaining_files, remaining_eps):
|
||||
matched_episodes.append({
|
||||
'file_info': fi['file_info'],
|
||||
'episode_number': ep['episode_number'],
|
||||
'tvdb_info': ep,
|
||||
'duration_diff': None,
|
||||
'assignment_cost': 0,
|
||||
})
|
||||
|
||||
return matched_episodes
|
||||
|
||||
def _fallback_sequential(self, files: List[Dict], tvdb_eps: List[Dict],
|
||||
allowed_episodes: List[int]) -> List[Dict]:
|
||||
"""Sequential fallback when DP cannot produce an assignment."""
|
||||
matched = []
|
||||
sorted_files = sorted(files, key=lambda x: x['filename'])
|
||||
allowed_set = set(allowed_episodes)
|
||||
avail = [ep for ep in tvdb_eps if ep['episode_number'] in allowed_set]
|
||||
for i, fi in enumerate(sorted_files):
|
||||
if i < len(avail):
|
||||
ep = avail[i]
|
||||
matched.append({
|
||||
'file_info': fi['file_info'],
|
||||
'episode_number': ep['episode_number'],
|
||||
'tvdb_info': ep,
|
||||
'duration_diff': None,
|
||||
})
|
||||
return matched
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Loading…
x
Reference in New Issue
Block a user