Số mảng con có trung vị bằng M trong hoán vị N số tự nhiên đầu tiên trong Python

Cho mảng A là hoán vị của N số tự nhiên đầu tiên và số M (M ≤ N). Ta cần đếm số mảng con mà trung vị của dãy bằng M. Trung vị là phần tử giữa sau khi sắp xếp tăng dần; nếu độ dài chẵn thì lấy phần tử bên trái trong hai phần tử giữa.

Ví dụ: A = [3, 5, 6, 4, 2], M = 5 -> kết quả 4 (các mảng con: [3,5,6], [5], [5,6], [5,6,4]).

Thuật toán giải quyết

• Duyệt mảng, giữ biến add (tăng khi > M, giảm khi < M).
• Dùng my_map lưu tần suất các giá trị add trước khi gặp M.
• Khi đã gặp M (has = True), với mỗi add hiện tại, kết quả cộng thêm my_map[add] và my_map[add-1].
• Trả về kết quả.

Ví dụ minh họa

def solve(arr, m):
    n = len(arr)
    my_map = {0: 1}
    has = False
    add = 0
    result = 0
    for i in range(n):
        if arr[i] < m:
            add -= 1
        elif arr[i] > m:
            add += 1
        if arr[i] == m:
            has = True
        if has:
            if add in my_map:
                result += my_map[add]
            if add - 1 in my_map:
                result += my_map[add - 1]
        else:
            my_map[add] = my_map.get(add, 0) + 1
    return result

arr = [3, 5, 6, 4, 2]
m = 5
print(solve(arr, m))  # 4

Đầu vào

[3, 5, 6, 4, 2], 5

Đầu ra

4