[python] 백준 16978번: 수열과 쿼리 22

링크: https://www.acmicpc.net/problem/16978

난이도: P4

---

문제

---

길이가 N인 수열 A1, A2, ..., AN이 주어진다. 이때, 다음 쿼리를 수행하는 프로그램을 작성하시오.

1 i v: Ai = v로 변경한다.
2 k i j: k번째 1번 쿼리까지 적용되었을 때, Ai, Ai+1, ..., Aj의 합을 출력한다.

 

입력

---

첫째 줄에 수열의 크기 N (1 ≤ N ≤ 100,000)이 주어진다.

둘째 줄에는 A1, A2, ..., AN이 주어진다. (1 ≤ Ai ≤ 1,000,000)

셋째 줄에는 쿼리의 개수 M (1 ≤ M ≤ 100,000)이 주어진다.

넷째 줄부터 M개의 줄에는 쿼리가 한 줄에 하나씩 주어진다. 1번 쿼리의 경우 1 ≤ i ≤ N, 1 ≤ v ≤ 1,000,000 이고, 2번 쿼리의 경우 1 ≤ i ≤ j ≤ N이고, 0 ≤ k ≤ (쿼리가 주어진 시점까지 있었던 1번 쿼리의 수)이다.

입력으로 주어지는 모든 수는 정수이다.

 

출력

---

모든 2번 쿼리마다 합을 출력한다.

 

예제 입력

5
1 2 3 4 5
7
1 2 5
2 0 1 3
2 1 1 3
1 4 2
2 0 2 5
2 1 2 5
2 2 2 5

예제 출력

6
9
14
17
15

풀이

---

기본적인 구간합 세그먼트 트리에서 쿼리에 변형이 추가된 문제다.

주어진 쿼리 순서대로 하면 시간초과가 발생할 수 밖에 없다.

우선, 1번쿼리를 add_q, 2번쿼리를 k_q로 따로 저장하고(이때, 기존의 k_q 순서를 idx로 추가해준다)

k_q를 k 오름차순으로 정렬해준다.

그렇게 되면 1번쿼리를 중복으로 진행할 필요 없이 최대 add_q 길이만큼만 진행해 주면 된다.

이후, 현재 1번쿼리의 어디까지 update진행되었는지 v로 따로 기억해주고 add_q를 돌아주면서 query 및 update를 적용해주고 

다시 k_q의 idx기준으로 정렬해준 뒤, 결과값들만 출력해준다.

 

제출코드

import sys
input=sys.stdin.readline

def build(le,ri,node):
    if le==ri:
        tree[node]=li[le]
        return
    mid=(le+ri)//2
    build(le,mid,node*2)
    build(mid+1,ri,node*2+1)
    tree[node]=tree[node*2]+tree[node*2+1]

def update(le,ri,node,idx,val):
    if idx<le or ri<idx:
        return 0
    if le==ri==idx:
        tree[node]=val
        return
    mid=(le+ri)//2
    update(le,mid,node*2,idx,val)
    update(mid+1,ri,node*2+1,idx,val)
    tree[node]=tree[node*2]+tree[node*2+1]

def query(st,end,node,le,ri):
    if end<le or ri<st:
        return 0
    if st<=le and ri<=end:
        return tree[node]
    mid=(le+ri)//2
    return query(st,end,node*2,le,mid)+query(st,end,node*2+1,mid+1,ri)


N=int(input())
li=[0]+list(map(int,input().split()))
tree=[0]*(N*4)
M=int(input())
add_q,k_q=[],[]
for idx in range(M):
    tmp=list(map(int,input().split()))
    if len(tmp)==3:
        _,i,v=tmp
        add_q.append([i,v])
    else:
        _,k,i,j=tmp
        k_q.append([k,i,j,idx])
k_q.sort(key=lambda x:x[0])
build(0,N,1)

v=0
for now in range(len(k_q)):
    k,i,j,idx=k_q[now]
    if v<k:
        for tar in range(v,k):
            tar_idx,tar_val=add_q[tar]
            update(0,N,1,tar_idx,tar_val)
    tmp_val=query(i,j,1,0,N)
    k_q[now].append(tmp_val)
    v=k
k_q.sort(key=lambda x:x[3])
for i in k_q:
    print(i[4])

 

---

틀린 부분이 있거나 인용한 부분에 대해 문제가 있을 시
댓글로 알려주시면 감사하겠습니다.