Computer >> บทช่วยสอนคอมพิวเตอร์ >  >> การเขียนโปรแกรม >> C++

Subarray Sum Equals K ใน C++

โจทย์: Subarray Sum Equals K

ให้อาร์เรย์จำนวนเต็ม nums และจำนวนเต็ม k ต้องหา จำนวนซับอาร์เรย์ติดต่อกัน (contiguous subarrays) ที่มีผลรวมเท่ากับ k อย่างแน่นอน

ตัวอย่าง: nums = [1, 1, 1], k = 2

ซับอาร์เรย์ที่ผลรวม = 2: [1,1] (index 0-1), [1,1] (index 1-2) → ผลลัพธ์: 2

อัลกอริทึม: Prefix Sum + Hash Map

กำหนด prefix_sum[i] = ผลรวมของ nums[0..i-1] ผลรวมของซับอาร์เรย์ [l, r] = prefix_sum[r+1] - prefix_sum[l]

เราต้องการ prefix_sum[r+1] - prefix_sum[l] = k → prefix_sum[l] = prefix_sum[r+1] - k

  1. ใช้ unordered_map<int, int> count เก็บจำนวนครั้งที่เจอแต่ละ prefix sum
  2. เริ่มต้น count[0] = 1 (prefix sum ว่าง = 0)
  3. วนลูปคำนวณ curr_sum สะสม:
    • หากมี curr_sum - k ใน map → เพิ่ม count[curr_sum - k] เข้า answer
    • เพิ่ม count[curr_sum]++

ตัวอย่างโค้ด C++

#include <bits/stdc++.h>
using namespace std;

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        unordered_map<int, int> count;
        count[0] = 1; // prefix sum ว่าง
        
        int curr_sum = 0;
        int ans = 0;
        
        for (int num : nums) {
            curr_sum += num;
            if (count.find(curr_sum - k) != count.end()) {
                ans += count[curr_sum - k];
            }
            count[curr_sum]++;
        }
        return ans;
    }
};

int main() {
    Solution ob;
    vector<int> v = {1, 1, 1};
    cout << ob.subarraySum(v, 2) << endl;
    return 0;
}

Input

[1,1,1] 2

Output

2

ความซับซ้อน (Complexity)

  • เวลา: O(n) - วนลูป 1 รอบ, hash map operation O(1) เฉลี่ย
  • พื้นที่: O(n) - hash map เก็บ prefix sum สูงสุด n ค่า

แก้ไขจากโค้ดต้นฉบับ: โค้ดต้นฉบับมีบั๊กที่ sums[-temp]++ ควรเป็น sums[curr_sum]++ (บวก ไม่ใช่ ลบ) และการตรวจสอบ k - temp ควรเป็น curr_sum - k ได้แก้ไขแล้วในตัวอย่างนี้