코딩테스트 JS 세그먼트 트리: 단일 대입과 구간 합

Node.js 24 CommonJS용 전체 풀이입니다. segment-tree-problem.cjs로 저장하세요. 클래스 선언, solve 함수, 표준 입력 처리가 모두 들어 있어 가이드 파일을 별도로 가져오지 않습니다. 예시 입력을 input.json에 저장하고 터미널에서 실행합니다.

Windows PowerShell:
Get-Content -Raw -Encoding UTF8 .input.json | node .segment-tree-problem.cjs

macOS / Linux:
node segment-tree-problem.cjs < input.json

class SegmentTree {
  constructor(values, combine = (a, b) => a + b, identity = 0) {
    if (!Array.isArray(values) || values.length > 1_000_000) {
      throw new RangeError('길이 1,000,000 이하 배열이 필요합니다.');
    }
    if (typeof combine !== 'function') throw new TypeError('결합 함수가 필요합니다.');
    this.n = values.length;
    this.combine = combine;
    this.identity = identity;
    this.base = 1;
    while (this.base < this.n) this.base *= 2;
    this.tree = Array(this.base * 2).fill(identity);
    for (let i = 0; i < this.n; i++) this.tree[this.base + i] = values[i];
    for (let i = this.base - 1; i > 0; i--) {
      this.tree[i] = combine(this.tree[i * 2], this.tree[i * 2 + 1]);
    }
  }
  boundary(i) {
    if (!Number.isInteger(i) || i < 0 || i > this.n) throw new RangeError('잘못된 경계입니다.');
  }
  set(index, value) {
    this.boundary(index);
    if (index === this.n) throw new RangeError('갱신 인덱스는 n보다 작아야 합니다.');
    let p = this.base + index;
    this.tree[p] = value;
    while (p > 1) {
      p = Math.floor(p / 2);
      this.tree[p] = this.combine(this.tree[p * 2], this.tree[p * 2 + 1]);
    }
  }
  query(left, right) {
    this.boundary(left);
    this.boundary(right);
    if (left > right) throw new RangeError('left는 right 이하여야 합니다.');
    let l = left + this.base, r = right + this.base;
    let accLeft = this.identity, accRight = this.identity;
    while (l < r) {
      if (l % 2 === 1) accLeft = this.combine(accLeft, this.tree[l++]);
      if (r % 2 === 1) accRight = this.combine(this.tree[--r], accRight);
      l = Math.floor(l / 2);
      r = Math.floor(r / 2);
    }
    return this.combine(accLeft, accRight);
  }
}

function solve({ values, ops }) {
  const tree = new SegmentTree(values), out = [];
  for (const [type, a, b] of ops) {
    if (type === 'replace') tree.set(a, b);
    else if (type === 'total') out.push(tree.query(a, b));
    else throw new Error('알 수 없는 작업입니다.');
  }
  return out;
}

module.exports = { solve };

if (require.main === module) {
  const input = JSON.parse(require('node:fs').readFileSync(0, 'utf8'));
  console.log(JSON.stringify(solve(input)));
}

첫 total은 2+7+1+3=13입니다. replace 이후 원본은 [4,2,5,1,3]이므로 같은 구간은 11입니다. 마지막 조회는 빈 구간이므로 항등원 0을 출력합니다. solve는 세그먼트 트리를 한 번 만들고 replace를 set에, total을 query에 대응시킵니다. 새 예상치를 delta로 처리하지 않는 것이 Fenwick 증가 예제와 다른 점입니다.

base는 n 이상인 가장 작은 2의 거듭제곱입니다. 원본 values[i]는 tree[base+i]에 놓고 나머지 잎은 항등원으로 채웁니다. 각 내부 노드는 왼쪽 자식과 오른쪽 자식을 그 순서로 결합합니다. 따라서 루트는 패딩을 포함해도 실제 배열 전체와 같은 결과를 갖습니다.

조회는 l,r을 잎 위치로 옮긴 뒤 아직 처리하지 않은 [l,r) 영역을 좁힙니다. l이 오른쪽 자식이면 왼쪽 형제가 범위 밖이므로 tree[l]만 받아 l을 하나 늘립니다. r이 홀수이면 r−1이 범위 안의 마지막 왼쪽 자식이므로 먼저 r을 줄이고 그 노드를 받습니다. 나머지는 완전한 형제 쌍이라 부모로 올릴 수 있습니다.

빌드 O(n+1), q개 작업 최악 O(q log(n+1)), 트리 O(n+1), 결과 O(q) 공간입니다. 숫자 합이 O(1)인 계약에서의 분석입니다.

표본과 경계를 직접 확인하려면 다음 코드를 segment-tree-check.cjs로 저장하고 같은 폴더에서 node segment-tree-check.cjs를 실행하세요. assert가 맞으면 출력 없이 종료하고, 다르면 예외와 함께 실패합니다.

const assert = require('node:assert/strict');
const { solve } = require('./segment-tree-problem.cjs');

assert.deepEqual(
  solve({
  "values": [
    4,
    2,
    7,
    1,
    3
  ],
  "ops": [
    [
      "total",
      1,
      5
    ],
    [
      "replace",
      2,
      5
    ],
    [
      "total",
      1,
      5
    ],
    [
      "total",
      0,
      0
    ]
  ]
}),
  [13,11,0],
);
assert.deepEqual(
  solve({
  "values": [],
  "ops": []
}),
  [],
);