927.三等分.md

927.三等分

/*
 * @lc app=leetcode.cn id=927 lang=typescript
 *
 * [927] 三等分
 */

// @lc code=start

function threeEqualParts(arr: number[]): number[] {}
// @lc code=end

解法 1: 三指针

function threeEqualParts(arr: number[]): number[] {
  const sum = arr.reduce((a, b) => a + b, 0)
  if (sum % 3 !== 0) return [-1, -1]
  if (sum === 0) return [0, 2]
  const n = arr.length,
    res = []
  let cur = 0
  for (let i = 0, j = 0; i < n; i++) {
    if (arr[i]) cur++
    if (cur === sum / 3) {
      cur = 0
      res.push(i)
    }
  }

  for (let [i, j, k] = res.map(num => num + 1); k < n; i++, j++, k++) {
    if (arr[i] !== 0 || arr[j] !== 0) return [-1, -1]
  }
  let cnt = 0
  for (let [i, j, k] = res; i >= 0; i--, j--, k--) {
    if (arr[i] !== arr[k] || arr[j] !== arr[k] || arr[i] !== arr[j]) return [-1, -1]
    cnt += arr[i]
    if (cnt === sum / 3) break
  }
  res[0] += n - res[2] - 1
  res[1] += n - res[2]
  res.pop()
  return res
}

Case

test.each([
  { input: { arr: [1, 0, 1, 0, 1] }, output: [0, 3] },
  { input: { arr: [1, 1, 0, 1, 1] }, output: [-1, -1] },
  { input: { arr: [1, 1, 0, 0, 1] }, output: [0, 2] },
])('input: arr = $input.arr', ({ input: { arr }, output }) => {
  expect(threeEqualParts(arr)).toEqual(output)
})