import numpy as np def replace_nans_with_local_statistic(arr, window_size=3, method='mean', max_iterations=10): """ Replace NaN values in a 1D array with the local mean or median of surrounding values, including handling of consecutive NaN segments via iterative refinement. 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. - max_iterations: Number of iterations to try filling NaNs (for consecutive NaNs). 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) result = arr.copy() n = len(arr) half_window = window_size // 2 for iteration in range(max_iterations): modified = False # Track whether any NaNs were filled in this pass for i in range(n): if np.isnan(result[i]): # Create a local window around the NaN start = max(0, i - half_window) end = min(n, i + half_window + 1) window = result[start:end] valid_values = window[~np.isnan(window)] if len(valid_values) > 0: if method == 'mean': result[i] = np.mean(valid_values) else: result[i] = np.median(valid_values) modified = True if not modified: # No changes in this iteration → stop break return result