ข้อ 19 · LC724 Find Pivot Index (หา pivot index) 🟢

โจทย์ (LC724): กำหนด array จำนวนเต็ม nums ให้คำนวณหา pivot index (ดัชนีจุดสมดุล) ของ array นี้ pivot index คือ index ที่ผลรวมของตัวเลขทั้งหมดที่อยู่ทางซ้ายของ index นั้น (strictly) เท่ากับผลรวมของตัวเลขทั้งหมดที่อยู่ทางขวาของ index นั้น (strictly) ถ้า index อยู่ริมซ้ายสุดของ array ให้ถือว่าผลรวมฝั่งซ้ายเป็น 0 เพราะไม่มีสมาชิกอยู่ทางซ้ายเลย กรณีริมขวาสุดก็เช่นเดียวกัน ให้ return pivot index ที่อยู่ซ้ายสุด ถ้าไม่มี index แบบนี้เลยให้ return -1

แนวทาง — ต้องใช้อะไร & คิดยังไง

ใช้ไอเดีย prefix sum (ผลรวมสะสม) แต่ครั้งนี้เราสนใจผลรวมสองฝั่งของแต่ละ index วิธี naive (ตรงตัว) คือทุก index ก็บวกฝั่งซ้ายใหม่และฝั่งขวาใหม่ ซึ่งกลายเป็น O(n^2) ช้าเกินไปเมื่อ array ยาว

ทริกคือรู้ผลรวมทั้งหมด total ไว้ก่อน แล้ว accumulate (สะสม) เฉพาะฝั่งซ้าย left_sum ระหว่าง iterate (เดิน) เมื่ออยู่ที่ index i ผลรวมฝั่งขวาหาได้ทันทีจาก total - left_sum - nums[i] ไม่ต้องบวกฝั่งขวาใหม่เลย

1. คำนวณ total = ผลรวมทั้ง array (หนึ่งรอบ)
2. initialize (ตั้งค่าเริ่มต้น) left_sum = 0 เก็บผลรวมของตัวที่อยู่ทางซ้ายของ i
3. iterate (วน) ทีละ index i พร้อมค่า n = nums[i]: คำนวณ right_sum = total - left_sum - n
4. ถ้า left_sum == right_sum แปลว่าเจอ pivot return i ทันที (ซ้ายสุดโดยธรรมชาติ)
5. ถ้ายังไม่สมดุล บวก n เข้า left_sum แล้วไปตัวถัดไป
6. จบ loop ยังไม่เจอ return -1

ไล่ทีละสเต็ป

iterate nums = [1, 7, 3, 6, 5, 6] ซึ่ง total = 28

def pivot_index(nums):
    total = sum(nums)   # ผลรวมทั้งหมด (จ่าย O(n) ครั้งเดียว)
    left_sum = 0        # ผลรวมของตัวที่อยู่ทางซ้ายของ i
    for i, n in enumerate(nums):
        # ผลรวมขวา = total - ฝั่งซ้าย - ตัวปัจจุบัน
        right_sum = total - left_sum - n
        if left_sum == right_sum:  # เช็คก่อนบวก n
            return i               # เจอ pivot ซ้ายสุด
        left_sum += n   # เลื่อน i ไปขวา ตัวปัจจุบันจึงกลายเป็นฝั่งซ้าย
    return -1           # ไล่จนจบไม่เจอ

print(pivot_index([1, 7, 3, 6, 5, 6]))  # 3
print(pivot_index([1, 2, 3]))           # -1
print(pivot_index([2, 1, -1]))          # 0

3
-1
0