import numpy as np def replace_nans_with_local_statistic(arr, window_size=3, method='mean'): """ Replace NaN values in a 1D array with the local mean or median of surrounding values. Parameters: - arr: Input 1D array with NaN values (list or numpy array). - window_size: Size of the window to compute local statistic (must be odd). - method: 'mean' or 'median' to specify the replacement method. Returns: - Modified array with NaNs replaced by local statistics. """ if method not in ['mean', 'median']: raise ValueError("Method must be 'mean' or 'median'") if window_size % 2 == 0: raise ValueError("Window size must be an odd number") arr = np.array(arr, dtype=float) # Ensure it's a float array to support NaN result = arr.copy() half_window = window_size // 2 for i in range(len(arr)): if np.isnan(arr[i]): start = max(0, i - half_window) end = min(len(arr), i + half_window + 1) window = arr[start:end] valid_window = window[~np.isnan(window)] if len(valid_window) > 0: if method == 'mean': result[i] = np.mean(valid_window) else: result[i] = np.median(valid_window) else: result[i] = np.nan # No valid values in window, leave as NaN return result