Chương trình tính tổng khoảng cách di chuyển tất cả quả bóng đến từng vị trí Python

Cho mảng nhị phân nums (0: ô trống, 1: có bóng). Tìm mảng L cùng kích thước, trong đó L[i] là tổng khoảng cách di chuyển tất cả bóng đến vị trí i (khoảng cách từ j đến i là |j-i|).

Ví dụ: nums=[1,1,0,1] → [4,3,4,5].

Giải thuật (Prefix sum hai chiều O(n))

1. Đếm right_count, right_sum (số bóng và tổng chỉ số ở bên phải)
2. Duyệt từ trái sang phải, cập nhật left_count, left_sum, right_count, right_sum
3. result[i] = left_sum + right_sum

Ví dụ triển khai

def solve(nums):
    if not nums:
        return []
    left_count = right_count = 0
    left_sum = right_sum = 0
    result = []
    for idx, num in enumerate(nums):
        if num:
            right_count += 1
            right_sum += idx
    for idx, num in enumerate(nums):
        result.append(left_sum + right_sum)
        if num:
            right_count -= 1
            left_count += 1
        left_sum += left_count
        right_sum -= right_count
    return result

print(solve([1, 1, 0, 1]))

Đầu vào

[1, 1, 0, 1]

Đầu ra

[4, 3, 4, 5]